Update app.py

app.py

@@ -1,56 +1,57 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import json
 
-# Load MagicoderS-CL-7B model
-MODEL_NAME = "ise-uiuc/Magicoder-S-CL-7B"
+MODEL_NAME = "tiiuae/falcon-rw-1b"  # lightweight and fast
+
 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
+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"
     )
-    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+    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)
 
-    # 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()
+        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()