init

app.py

@@ -1,6 +1,7 @@
 import spaces
 import gradio as gr
 from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM, TrainerCallback
+from transformers import DataCollatorForSeq2Seq
 from datasets import load_dataset
 import traceback
 from huggingface_hub import login
@@ -41,17 +42,30 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
         def tokenize_function(examples):
             max_length = 16
             # Assuming 'text' is the input and 'target' is the expected output
-            model_inputs = tokenizer(examples['text'], max_length=max_length, truncation=True)
+            model_inputs = tokenizer(
+                examples['text'], 
+                max_length=None,  # Set to None for dynamic padding
+                padding=False,     # Disable padding here, we will handle it later
+                truncation=True,
+                return_tensors="pt"  # Return PyTorch tensors
+            )
         
             # Setup the decoder input IDs (shifted right)
             with tokenizer.as_target_tokenizer():
-                labels = tokenizer(examples['target'], max_length=max_length, truncation=True)
+                labels = tokenizer(
+                    examples['target'], 
+                    max_length=None,  # Set to None for dynamic padding
+                    padding=False,     # Disable padding here, we will handle it later
+                    truncation=True,
+                    return_tensors="pt"                    
+                )
         
             # Add labels to the model inputs
             model_inputs["labels"] = labels["input_ids"]
             return model_inputs
     
         tokenized_datasets = dataset.map(tokenize_function, batched=True)
+        data_collator = DataCollatorForSeq2Seq(tokenizer, model=model)
     
         # Set training arguments
         training_args = TrainingArguments(
@@ -63,7 +77,7 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
             per_device_eval_batch_size=1, 
             num_train_epochs=int(num_epochs),
             weight_decay=0.01,
-            #gradient_accumulation_steps=grad*0.1,
+            gradient_accumulation_steps=int(grad),
             load_best_model_at_end=True,
             metric_for_best_model="accuracy",
             greater_is_better=True,
@@ -79,8 +93,8 @@ def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch
         trainer = Trainer(
             model=model,
             args=training_args,
-            train_dataset=tokenized_datasets['train'],
-            eval_dataset=tokenized_datasets['test'],
+            train_dataset=data_collator['train'],
+            eval_dataset=data_collator['test'],
             #callbacks=[LoggingCallback()], 
         )