Update app.py

app.py

@@ -28,8 +28,7 @@ except ImportError as e:
     from accelerate import Accelerator
     import bitsandbytes
 
-# Rest of your script remains the same...
-# Model setup, training function, Gradio UI, etc., as shown in your previous script
+# Model setup
 MODEL_ID = "meta-llama/Llama-2-7b-hf"  # Use Llama-2-7b; switch to "meta-llama/Llama-3-8b-hf" for Llama 3
 tokenizer = LlamaTokenizer.from_pretrained(MODEL_ID)
 
@@ -122,39 +121,4 @@ def train_ui_tars(file):
         def custom_data_collator(features):
             batch = {
                 "input_ids": torch.stack([f["input_ids"] for f in features]),
-                "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()
+                "attention_mask": torch.stack([f["