Update app.py

app.py

@@ -1,46 +1,16 @@
-import gradio as gr
-from pathlib import Path
 import os
-import zipfile
-
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-
-model_dir = Path(__file__).parent / "trained_model"
-
-# Gracefully fail if model not trained yet
-if not model_dir.exists():
-    raise RuntimeError("❌ Model not trained yet! Please train it using train.py.")
-
-# Load locally
-tokenizer = AutoTokenizer.from_pretrained(model_dir, local_files_only=True)
-model = AutoModelForCausalLM.from_pretrained(model_dir, local_files_only=True)
-
-def generate_code(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
-    outputs = model.generate(**inputs, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+import gradio as gr
 
-def zip_model():
-    zip_path = Path("trained_model.zip")
-    with zipfile.ZipFile(zip_path, "w") as zf:
-        for foldername, _, filenames in os.walk(model_dir):
-            for file in filenames:
-                filepath = Path(foldername) / file
-                zf.write(filepath, filepath.relative_to(model_dir.parent))
-    return str(zip_path)
+# Check for trained model zip
+model_zip = "trained_model.zip"
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 Trained Python Code Generator")
-
-    prompt = gr.Textbox(label="Enter prompt")
-    output = gr.Textbox(label="AI Response")
-
-    with gr.Row():
-        gen_btn = gr.Button("Generate")
-        dl_btn = gr.Button("Download Model ZIP")
+    gr.Markdown("## 🤖 Python AI Model Trainer")
 
-    gen_btn.click(generate_code, inputs=prompt, outputs=output)
-    dl_btn.click(zip_model, outputs=gr.File(label="Download ZIP"))
+    if os.path.exists(model_zip):
+        gr.Markdown("### ✅ Model trained! Click below to download:")
+        gr.File(value=model_zip, label="Download trained model")
+    else:
+        gr.Markdown("### 🚫 No trained model found yet. Run train.py to train it.")
 
 demo.launch()