File size: 1,612 Bytes
a31c67b
1b40e2f
41e5367
2ea1871
2223fe4
1b40e2f
f7ce9eb
6f6d780
1b40e2f
2223fe4
41e5367
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2ea1871
41e5367
2223fe4
41e5367
2ea1871
41e5367
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
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()