Update app.py

app.py

@@ -1,50 +1,46 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
 from pathlib import Path
-import torch
-import zipfile
 import os
+import zipfile
 
-# Path to local model directory
-model_dir = Path(__file__).parent / "trained_model"
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+model_dir = Path(_file_).parent / "trained_model"
 
-# Load model and tokenizer with local_files_only=True
+# 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)
 
-# Inference function
 def generate_code(prompt):
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=128,
-            pad_token_id=tokenizer.eos_token_id
-        )
+    outputs = model.generate(**inputs, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# Create a downloadable zip of the trained model
 def zip_model():
     zip_path = Path("trained_model.zip")
-    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
-        for root, _, files in os.walk(model_dir):
-            for file in files:
-                filepath = Path(root) / file
-                arcname = filepath.relative_to(model_dir.parent)
-                zipf.write(filepath, arcname)
+    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)
 
-# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 🧠 Trained Python Code Generator")
-    prompt = gr.Textbox(label="Enter a prompt for code completion")
+    gr.Markdown("# 🤖 Trained Python Code Generator")
+
+    prompt = gr.Textbox(label="Enter prompt")
     output = gr.Textbox(label="AI Response")
 
     with gr.Row():
-        generate_btn = gr.Button("Generate")
-        download_btn = gr.Button("Download Trained Model")
+        gen_btn = gr.Button("Generate")
+        dl_btn = gr.Button("Download Model ZIP")
 
-    generate_btn.click(fn=generate_code, inputs=prompt, outputs=output)
-    download_btn.click(fn=zip_model, outputs=gr.File(label="Download ZIP"))
+    gen_btn.click(generate_code, inputs=prompt, outputs=output)
+    dl_btn.click(zip_model, outputs=gr.File(label="Download ZIP"))
 
 demo.launch()