Update app.py

app.py

@@ -1,32 +1,85 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import time
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+# ======================================
+# Device setup
+# ======================================
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-device = torch.device('cuda')
+# ======================================
+# Model config
+# ======================================
+MODEL_NAME = "Salesforce/codegen-350M-mono"
 
-model_name = "Salesforce/codegen-350M-mono"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+# Required for GPT-style models
+tokenizer.pad_token = tokenizer.eos_token
+model.config.pad_token_id = tokenizer.eos_token_id
 
-model.to(device)
+model.to(DEVICE)
+model.eval()
 
-def generate_code(prompt):
-    full_prompt = f"# HTML/CSS/JS website code for:\n# {prompt}\n"
+# ======================================
+# Generation function
+# ======================================
+def generate_code(prompt: str) -> str:
+    if not prompt.strip():
+        return ""
+
+    # Easter egg
     if prompt == "OzLabsWereHere":
         time.sleep(10)
         return "M'en va t'à la fontaine pour y pêcher du poisson, la Ziguezon zinzon"
-    inputs = tokenizer(full_prompt, return_tensors="pt").input_ids
-    outputs = model.generate(inputs, max_length=1024, temperature=0.7)
-    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    html_start = decoded.find('<')
+
+    full_prompt = (
+        "# Generate a complete HTML/CSS/JS website.\n"
+        "# Return ONLY valid code.\n\n"
+        f"{prompt}\n"
+    )
+
+    inputs = tokenizer(
+        full_prompt,
+        return_tensors="pt",
+        padding=True,
+        truncation=True
+    )
+
+    # 🔑 Move inputs to the same device as the model
+    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs,
+            max_length=1024,
+            temperature=0.7,
+            do_sample=True,
+            top_p=0.95
+        )
+
+    decoded = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+    # Extract HTML if possible
+    html_start = decoded.find("<")
     return decoded[html_start:] if html_start != -1 else decoded
 
-gr.Interface(
+
+# ======================================
+# Gradio UI
+# ======================================
+app = gr.Interface(
     fn=generate_code,
-    inputs=gr.Textbox(lines=6, label="Describe your website idea"),
-    outputs=gr.Code(label="Generated HTML/CSS/JS"),
+    inputs=gr.Textbox(
+        lines=6,
+        placeholder="Describe the website you want to generate...",
+        label="Website description"
+    ),
+    outputs=gr.Code(label="Generated HTML / CSS / JS"),
     title="Oz AI Website Generator",
-    description="Describe your idea in English or Hindi. This will return only HTML/CSS/JS."
-).launch()
+    description="Describe a website idea. The model returns only HTML/CSS/JS."
+)
+
+app.launch()