add tests - data loader test

easy_tpp/preprocess/dataset.py

@@ -116,7 +116,7 @@ def get_dt_stats(self):
             s_2_y = dts.var()
             m = dts.shape[0]
             n += m
-            # Formulat taken from https://math.stackexchange.com/questions/3604607/can-i-work-out-the-variance-in-batches
+            # Formula taken from https://math.stackexchange.com/questions/3604607/can-i-work-out-the-variance-in-batches
             s_2_x = (((n - 1) * s_2_x + (m - 1) * s_2_y) / (n + m - 1)) + (
                         (n * m * ((x_bar - y_bar) ** 2)) / ((n + m) * (n + m - 1)))
             x_bar = (n * x_bar + m * y_bar) / (n + m)

easy_tpp/preprocess/event_tokenizer.py

@@ -414,7 +414,9 @@ def _pad(
                                                                              max_len=max_length,
                                                                              dtype=np.int64)
         else:
-            batch_output = encoded_inputs
+            batch_output[self.model_input_names[0]] = np.array(encoded_inputs[self.model_input_names[0]])
+            batch_output[self.model_input_names[1]] = np.array(encoded_inputs[self.model_input_names[1]])
+            batch_output[self.model_input_names[2]] = np.array(encoded_inputs[self.model_input_names[2]])
 
         # non_pad_mask; replaced the use of event types by using the original sequence length
         seq_pad_mask = np.full_like(batch_output[self.model_input_names[2]], fill_value=True, dtype=bool)

examples/hf_data_loader.py

@@ -17,4 +17,5 @@ def load_data_from_hf(hf_dir=None, local_dir=None):
 if __name__ == '__main__':
     # in case one fails to load from hf directly
     # one can load the json data file locally
-    load_data_from_hf(hf_dir=None, local_dir={'validation':'dev.json'})
+    # load_data_from_hf(hf_dir=None, local_dir={'validation':'dev.json'})
+    load_data_from_hf(hf_dir='easytpp/taxi')