do not use sample packing for running predicitons

Author: ChaoPang

src/cehrbert/data_generators/hf_data_generator/hf_dataset_collator.py

@@ -116,6 +116,10 @@ def __call__(self, examples):
             )
             # Set the visit_segments of the CLS token to a default value 0 because this doesn't belong to a visit
             batch["visit_segments"] = torch.cat([torch.full((batch_size, 1), 0), batch["visit_segments"]], dim=1)
+        else:
+            assert (
+                    batch["attention_mask"].shape[0] == 1
+            ), f"batch['attention_mask'].shape[0] must be 0 in sample packing"
 
         # This is the most crucial logic for generating the training labels
         if self.is_pretraining:

src/cehrbert/runners/hf_cehrbert_finetune_runner.py

@@ -330,7 +330,12 @@ def main():
             dataset=processed_dataset["test"],
             batch_size=per_device_eval_batch_size,
             num_workers=training_args.dataloader_num_workers,
-            collate_fn=data_collator,
+            collate_fn=CehrBertDataCollator(
+                tokenizer=tokenizer,
+                max_length=config.max_position_embeddings,
+                is_pretraining=False,
+                mlm_probability=config.mlm_probability,
+            ),
             pin_memory=training_args.dataloader_pin_memory,
         )
         do_predict(test_dataloader, model_args, training_args)