app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import json

# Load MagicoderS-CL-7B model
MODEL_NAME = "ise-uiuc/Magicoder-S-CL-7B"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# System prompt
SYSTEM_INSTRUCTION = (
    "You are a helpful assistant that writes UI code. "
    "Respond only in JSON format with the following keys: "
    "\"filename\", \"html\", \"css\", and \"js\". "
    "Do not include explanations. Just return a JSON object."
)

def generate_code(user_prompt):
    prompt = f"{SYSTEM_INSTRUCTION}\nPrompt: {user_prompt}"

    inputs = tokenizer(prompt, return_tensors="pt", truncation=True).to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=512,
        temperature=0.7,
        top_p=0.9,
        do_sample=True
    )
    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()

    # Try to parse JSON
    try:
        result = json.loads(decoded)
    except json.JSONDecodeError:
        # fallback if output is not valid JSON
        result = {
            "filename": "index.html",
            "html": decoded,
            "css": "",
            "js": ""
        }

    return result

# Gradio Interface
iface = gr.Interface(
    fn=generate_code,
    inputs=gr.Textbox(lines=2, placeholder="Describe the UI you want..."),
    outputs="json",
    title="UI Code Generator",
    description="Enter a prompt to generate HTML, CSS, and JS code in structured JSON format."
)

if __name__ == "__main__":
    iface.launch()