Update app.py

app.py

@@ -117,8 +117,43 @@ def train_ui_tars(file):
             warmup_steps=100,
         )
        
-        # Custom data collator for Llama
+        # Custom data collator for Llama (corrected)
         def custom_data_collator(features):
             batch = {
                 "input_ids": torch.stack([f["input_ids"] for f in features]),
-                "attention_mask": torch.stack([f["
+                "attention_mask": torch.stack([f["attention_mask"] for f in features]),
+                "labels": torch.stack([f["labels"] for f in features]),
+            }
+            return batch
+
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=tokenized_dataset,
+            data_collator=custom_data_collator,
+        )
+       
+        # Step 4: Start training
+        trainer.train()
+       
+        # Step 5: Save the model
+        model.save_pretrained("./fine_tuned_llama")
+        tokenizer.save_pretrained("./fine_tuned_llama")
+       
+        return "Training completed successfully! Model saved to ./fine_tuned_llama"
+   
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Gradio UI
+with gr.Blocks(title="Model Fine-Tuning Interface") as demo:
+    gr.Markdown("# Llama Fraud Detection Fine-Tuning UI")
+    gr.Markdown("Upload a JSON file with 'input' and 'output' pairs to fine-tune the Llama model on your fraud dataset.")
+   
+    file_input = gr.File(label="Upload Fraud Dataset (JSON)")
+    train_button = gr.Button("Start Fine-Tuning")
+    output = gr.Textbox(label="Training Status")
+   
+    train_button.click(fn=train_ui_tars, inputs=file_input, outputs=output)
+
+demo.launch()