Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from datasets import load_dataset
-from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
+from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForSeq2Seq
 
 # Load the dataset
 dataset = load_dataset("json", data_files="dataset.jsonl")
@@ -12,29 +12,48 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Tokenize the dataset
 def tokenize_function(examples):
-    return tokenizer(examples["input"], text_target=examples["output"], truncation=True)
+    return tokenizer(
+        examples["input"],
+        text_target=examples["output"],
+        truncation=True,  # Truncate sequences longer than max_length
+        max_length=512,   # Adjust this based on your use case
+        padding="max_length"  # Pad shorter sequences to max_length
+    )
 
 tokenized_dataset = dataset.map(tokenize_function, batched=True)
 
+for i, example in enumerate(tokenized_dataset["train"]):
+    input_len = len(example["input_ids"])
+    output_len = len(example["labels"])
+    print(f"Example {i}: Input length = {input_len}, Output length = {output_len}")
+    
 # Define training arguments
 training_args = TrainingArguments(
     output_dir="./results",
-    overwrite_output_dir=True,
-    eval_strategy="epoch",  # Updated
-    learning_rate=5e-5,
-    per_device_train_batch_size=2,
+    per_device_train_batch_size=1,  # Smaller batch size
+    gradient_accumulation_steps=8,  # Accumulate gradients to simulate larger batch size
     num_train_epochs=3,
-    save_strategy="epoch",
     logging_dir="./logs",
-    logging_strategy="epoch",
+    logging_strategy="steps",
+    save_strategy="epoch",
+    eval_strategy="epoch",
+    learning_rate=5e-5,
+    overwrite_output_dir=True,
+)
+
+data_collator = DataCollatorForSeq2Seq(
+    tokenizer,
+    model=model,
+    padding=True,  # Enable dynamic padding
+    return_tensors="pt"
 )
 
-# Trainer setup
 trainer = Trainer(
     model=model,
     args=training_args,
     train_dataset=tokenized_dataset["train"],
     eval_dataset=tokenized_dataset["train"],
+    data_collator=data_collator,  # Use dynamic padding
 )
 
 # Train the model