Update app.py

app.py

@@ -1,17 +1,37 @@
-import os
-import subprocess
 import gradio as gr
+import os
+import shutil
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Load the trained model and tokenizer
+model_dir = "trained_model"
+model = AutoModelForCausalLM.from_pretrained(model_dir)
+tokenizer = AutoTokenizer.from_pretrained(model_dir)
+
+# Generate response function
+def generate_code(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(**inputs, max_length=256)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# Save model to accessible directory for download
+def save_model_for_download():
+    if os.path.exists("/home/user/app/files/trained_model"):
+        shutil.rmtree("/home/user/app/files/trained_model")
+    shutil.copytree("trained_model", "/home/user/app/files/trained_model")
+    return "Model saved to files tab. You can now download it."
 
-def run_training():
-    result = subprocess.run(["python3", "train.py"], capture_output=True, text=True)
-    return result.stdout + "\n\n" + result.stderr
+with gr.Blocks() as demo:
+    gr.Markdown("# Trained Python Code Generator")
+    with gr.Row():
+        inp = gr.Textbox(label="Enter a prompt")
+        out = gr.Textbox(label="Generated Python Code")
+    btn = gr.Button("Generate")
+    btn.click(fn=generate_code, inputs=inp, outputs=out)
 
-iface = gr.Interface(
-    fn=run_training,
-    inputs=[],
-    outputs="text",
-    title="Model Trainer",
-    description="Click the button to train the model using train.py"
-)
+    gr.Markdown("## Download Trained Model")
+    save_btn = gr.Button("Save model to files tab")
+    save_status = gr.Textbox(label="Status")
+    save_btn.click(fn=save_model_for_download, outputs=save_status)
 
-iface.launch()
+demo.launch(