Update app.py

app.py

@@ -1,30 +1,49 @@
+import os
+import subprocess
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
-
-model_name = "Percy3822/python_ai_coder"  # Replace with your model repo name after training
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
-
-def generate_code(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=128,
-            do_sample=True,
-            temperature=0.8,
-            top_p=0.95,
-            eos_token_id=tokenizer.eos_token_id
-        )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-iface = gr.Interface(
-    fn=generate_code,
-    inputs=gr.Textbox(lines=5, placeholder="Ask me to write/fix/explain Python code..."),
-    outputs="text",
-    title="Python AI Assistant (Trained on StarCoder)",
-    description="Ask it to write functions, fix bugs, explain code, etc."
-)
-
-iface.launch()
+import threading
+import time
+import zipfile
+
+output_path = "train_output"
+zipped_model = "python_ai_trained_model.zip"
+status = gr.Textbox(label="Status", value="Ready", interactive=False)
+download_link = gr.File(label="Download Trained Model", visible=False)
+
+def run_training():
+    global zipped_model
+    status.value = "Training started... this may take a while (15–60+ mins)."
+
+    # Run train.py
+    subprocess.run(["python", "train.py"], check=True)
+
+    # Compress trained model
+    status.value = "Training complete. Compressing model..."
+    with zipfile.ZipFile(zipped_model, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        for root, _, files in os.walk(output_path):
+            for file in files:
+                filepath = os.path.join(root, file)
+                arcname = os.path.relpath(filepath, output_path)
+                zipf.write(filepath, arcname)
+
+    # Move zip to Gradio-visible path
+    status.value = "Done. Model ready for download."
+    download_link.visible = True
+
+def start_training():
+    training_thread = threading.Thread(target=run_training)
+    training_thread.start()
+    return "Training in progress...", None
+
+with gr.Blocks() as demo:
+    gr.Markdown("## 🧠 Python AI Model Trainer (StarCoder 7B)")
+    gr.Markdown("Click the button below to start fine-tuning your custom Python AI. After training, download the model and use it anywhere.")
+    
+    with gr.Row():
+        train_btn = gr.Button("🚀 Start Training")
+        train_btn.click(start_training, outputs=[status, download_link])
+
+    status.render()
+    download_link.render()
+
+demo.launch()