Update app.py

app.py

@@ -7,8 +7,6 @@ from threading import Thread
 hf_token = os.getenv("HF_TOKEN")
 model_id = "ZyperAI/Z-AI-0.1-1.1B-Code.web"
 
-# Professional loading logic: Since Gradio 6.x runs as a persistent server, 
-# global variables are naturally 'cached' for the duration of the process.
 print("Loading model and tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(
     model_id, 
@@ -24,19 +22,31 @@ model = AutoModelForCausalLM.from_pretrained(
 print("Model loaded successfully.")
 
 def generate_code(prompt, history):
-    messages = history + [{"role": "user", "content": prompt}]
+    # Fix 1: Properly structure history for Gradio 6's list-of-dicts style
+    messages = []
+    for msg in history:
+        # Prevent appending empty or broken dictionary structures
+        if msg.get("content"):
+            messages.append({"role": msg["role"], "content": msg["content"]})
+            
+    messages.append({"role": "user", "content": prompt})
     
-    # Prepare inputs using the model's chat template
-    inputs = tokenizer.apply_chat_template(
-        messages, 
-        add_generation_prompt=True, 
-        return_tensors="pt"
-    ).to("cpu")
+    # Fix 2: Explicitly handle chat template errors if tokens are missing
+    try:
+        inputs = tokenizer.apply_chat_template(
+            messages, 
+            add_generation_prompt=True, 
+            return_tensors="pt"
+        ).to("cpu")
+    except Exception:
+        # Fallback if the specific model lacks a pre-configured chat template
+        fallback_prompt = "\n".join([f"{m['role']}: {m['content']}" for m in messages]) + "\nassistant:"
+        inputs = tokenizer(fallback_prompt, return_tensors="pt").input_ids.to("cpu")
     
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
     generation_kwargs = dict(
-        inputs=inputs,
+        input_ids=inputs, # Fix 3: Transformers generation kwargs expects 'input_ids', not 'inputs'
         streamer=streamer,
         max_new_tokens=1024,
         do_sample=True,
@@ -47,16 +57,19 @@ def generate_code(prompt, history):
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
     
-    response = ""
+    # Fix 4: Gradio 6 gr.Chatbot (type="messages") yields back the full history list, 
+    # not just a single raw string.
+    updated_history = messages.copy()
+    updated_history.append({"role": "assistant", "content": ""})
+    
     for new_text in streamer:
-        response += new_text
-        yield response
+        updated_history[-1]["content"] += new_text
+        yield updated_history
 
 # Gradio 6.x UI setup
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
     gr.Markdown("# ⚡ **Z-AI Web Coder**")
     
-    # In Gradio 6, type="messages" is the standard for the chatbot component
     chatbot = gr.Chatbot(height=500, show_copy_button=True, type="messages")
     with gr.Row():
         msg = gr.Textbox(
@@ -65,10 +78,20 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
             scale=9
         )
         submit = gr.Button("Build", variant="primary", scale=1)
-    msg.submit(generate_code, [msg, chatbot], [chatbot])
-    submit.click(generate_code, [msg, chatbot], [chatbot])
-    msg.submit(lambda: "", None, [msg])
-    submit.click(lambda: "", None, [msg])
+    
+    # Fix 5: Use a unified event pipeline so input clearing 
+    # doesn't disrupt the streaming text generator.
+    submit_click = submit.click(
+        generate_code, 
+        inputs=[msg, chatbot], 
+        outputs=[chatbot]
+    ).then(lambda: "", None, [msg])
+    
+    msg_submit = msg.submit(
+        generate_code, 
+        inputs=[msg, chatbot], 
+        outputs=[chatbot]
+    ).then(lambda: "", None, [msg])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()