Fix bug when TrainingArguments and Trainer

spanish_medica_llm.py

@@ -16,10 +16,10 @@ from datasets import load_dataset, concatenate_datasets
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
-    BitsAndBytesConfig,
+    BitsAndBytesConfig, 
+    DataCollatorForLanguageModeling,
     TrainingArguments,
-    Trainer,
-    DataCollatorForLanguageModeling
+    Trainer
 )
 
 from accelerate import FullyShardedDataParallelPlugin, Accelerator
@@ -399,7 +399,7 @@ def getTokenizedDataset(dataset, tokenizer):
         return dataset
 
     return  dataset.map(
-        lambda element : tokenize(element, tokenizer) ,
+        lambda element : tokenize(element, tokenizer),
         batched = True,
         remove_columns = dataset["train"].column_names
     )
@@ -497,8 +497,7 @@ def configAndRunTraining(basemodel, dataset, eval_dataset, tokenizer):
     else:
         tokenizer.pad_token = tokenizer.eos_token
         data_collator_pretrain = DataCollatorForLanguageModeling(tokenizer, mlm = False)
-
-        training_args = transformers.TrainingArguments(
+        training_args = TrainingArguments(
                 output_dir=output_dir,
                 push_to_hub = True,
                 hub_private_repo = False,
@@ -524,7 +523,7 @@ def configAndRunTraining(basemodel, dataset, eval_dataset, tokenizer):
                 bf16=False
             )
         
-        trainer = transformers.Trainer(
+        trainer = Trainer(
                      model= basemodel,
                      train_dataset = dataset['train'],
                      eval_dataset = eval_dataset,