fixed a bug where the sample packing was never used in do_predict

Author: ChaoPang

src/cehrbert/runners/hf_cehrbert_finetune_runner.py

@@ -19,7 +19,7 @@
 from tqdm import tqdm
 from transformers import EarlyStoppingCallback, Trainer, TrainingArguments, set_seed
 from transformers.trainer_utils import is_main_process
-from transformers.utils import logging
+from transformers.utils import is_flash_attn_2_available, logging
 
 from cehrbert.data_generators.hf_data_generator.cache_util import CacheFileCollector
 from cehrbert.data_generators.hf_data_generator.hf_dataset import create_cehrbert_finetuning_dataset
@@ -168,7 +168,9 @@ def load_finetuned_model(model_args: ModelArguments, model_name_or_path: str) ->
     torch_dtype = get_torch_dtype(model_args.torch_dtype)
     try:
         model = finetune_model_cls.from_pretrained(
-            model_name_or_path, torch_dtype=torch_dtype, attn_implementation=model_args.attn_implementation
+            model_name_or_path,
+            torch_dtype=torch_dtype,
+            attn_implementation=("flash_attention_2" if is_flash_attn_2_available() else "eager"),
         )
         if torch_dtype == torch.bfloat16:
             return model.bfloat16()
@@ -336,12 +338,7 @@ def main():
             dataset=processed_dataset["test"],
             batch_size=per_device_eval_batch_size,
             num_workers=training_args.dataloader_num_workers,
-            collate_fn=CehrBertDataCollator(
-                tokenizer=tokenizer,
-                max_length=config.max_position_embeddings,
-                is_pretraining=False,
-                mlm_probability=config.mlm_probability,
-            ),
+            collate_fn=data_collator,
             pin_memory=training_args.dataloader_pin_memory,
         )
         do_predict(test_dataloader, model_args, training_args)