Update app.py

app.py

@@ -45,29 +45,33 @@ else:
     )
     model = get_peft_model(model, lora_config)
 
-    # Load dataset and tokenize
+    # Load dataset
     dataset = load_dataset("mbpp", split="train")
 
+    # 🔥 Fix: Set `labels` properly
     def tokenize_function(examples):
-        return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)
+        inputs = tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)
+        inputs["labels"] = inputs["input_ids"].copy()  # ✅ Fix: Ensure labels exist
+        return inputs
 
     tokenized_dataset = dataset.map(tokenize_function, batched=True)
 
     # Data collator
     data_collator = DataCollatorForSeq2Seq(tokenizer, return_tensors="pt")
 
+    # Training arguments
     training_args = TrainingArguments(
-    per_device_train_batch_size=1,
-    num_train_epochs=1,  # Reduce epochs for quick training
-    learning_rate=3e-4,
-    output_dir=model_path,
-    save_strategy="epoch",
-    logging_dir="./logs",
-    logging_steps=10,
-    save_total_limit=2,
-    evaluation_strategy="no",  # ✅ No evaluation dataset needed
-    load_best_model_at_end=False  # ❌ Fix: Remove this to avoid error
-)
+        per_device_train_batch_size=1,
+        num_train_epochs=1,  # Reduce epochs for quick training
+        learning_rate=3e-4,
+        output_dir=model_path,
+        save_strategy="epoch",
+        logging_dir="./logs",
+        logging_steps=10,
+        save_total_limit=2,
+        evaluation_strategy="no",  # ✅ Fix: No evaluation dataset needed
+        load_best_model_at_end=False  # ✅ Fix: Prevents conflict with no eval dataset
+    )
 
     # Trainer
     trainer = Trainer(