app.py

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."

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)

    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)

demo.launch()