Update app.py

app.py

@@ -1,57 +1,56 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 import json
 
-MODEL_NAME = "google/flan-t5-base"
-
+# Load CodeT5 model
+MODEL_NAME = "Salesforce/codet5-base"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
 
-def generate_code(prompt, history):
-    history = history or []
-    system_instruction = (
-        "You are a coding assistant. "
-        "Respond only in valid JSON format like this:\n"
-        "{\n"
-        "  \"filename\": \"index.html\",\n"
-        "  \"html\": \"...HTML code...\",\n"
-        "  \"css\": \"...CSS code...\",\n"
-        "  \"js\": \"...JavaScript code...\"\n"
-        "}\n"
+# 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
     )
-    full_prompt = system_instruction + "\nPrompt: " + prompt
-    inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=512)
-    outputs = model.generate(**inputs, max_new_tokens=512)
-    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
 
+    # Try to parse JSON
     try:
-        parsed = json.loads(decoded)
-        html = parsed.get("html", "")
-        css = parsed.get("css", "")
-        js = parsed.get("js", "")
-    except Exception:
-        html, css, js = "", "", ""
-
-    history.append((prompt, decoded))
-    return html, css, js, history
-
-def clear_history():
-    return "", "", "", []
-
-with gr.Blocks() as demo:
-    chat_history = gr.State([])
-
-    with gr.Row():
-        inp = gr.Textbox(label="Prompt", lines=2)
-        send_btn = gr.Button("Generate")
-        clear_btn = gr.Button("New Chat")
-
-    with gr.Row():
-        html_out = gr.Code(label="HTML", language="html")
-        css_out = gr.Code(label="CSS", language="css")
-        js_out = gr.Code(label="JavaScript", language="javascript")
-
-    send_btn.click(generate_code, [inp, chat_history], [html_out, css_out, js_out, chat_history])
-    clear_btn.click(clear_history, outputs=[html_out, css_out, js_out])
-
-demo.launch()
+        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()