Update app.py

app.py

@@ -4,19 +4,27 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from threading import Thread
 
-# This pulls the token you just saved in the Space settings
 hf_token = os.getenv("HF_TOKEN")
-
 model_id = "ZyperAI/Z-AI-0.1-1.1B-Code.web"
 
-# We pass the token to both the tokenizer and the model
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id, 
-    torch_dtype=torch.float32, 
-    device_map="cpu",
-    token=hf_token
-)
+@gr.cache_resource
+def load_assets():
+    # Adding use_fast=False can bypass some backend instantiation errors 
+    # while still providing full functionality.
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_id, 
+        token=hf_token, 
+        use_fast=False 
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id, 
+        torch_dtype=torch.float32, 
+        device_map="cpu",
+        token=hf_token
+    )
+    return tokenizer, model
+
+tokenizer, model = load_assets()
 
 def generate_code(prompt, history):
     messages = []
@@ -25,7 +33,13 @@ def generate_code(prompt, history):
         messages.append({"role": "assistant", "content": assistant})
     messages.append({"role": "user", "content": prompt})
     
-    inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to("cpu")
+    # Apply the chat template specific to the ZyperAI model
+    inputs = tokenizer.apply_chat_template(
+        messages, 
+        add_generation_prompt=True, 
+        return_tensors="pt"
+    ).to("cpu")
+    
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
     generation_kwargs = dict(
@@ -45,33 +59,20 @@ def generate_code(prompt, history):
         response += new_text
         yield response
 
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue", neutral_hue="slate")) as demo:
-    gr.Markdown("# **Z-AI Web Coder**")
-    gr.Markdown("### *Professional AI Coding Assistant Powered by ZyperAI*")
+# Professional Dark-Themed UI
+with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", neutral_hue="zinc")) as demo:
+    gr.Markdown("# ⚡ **Z-AI Web Coder**")
+    gr.Markdown("Optimized 1.1B model for high-speed web development.")
     
-    chatbot = gr.Chatbot(
-        height=600, 
-        show_copy_button=True, 
-        bubble_full_width=False,
-        type="messages"
-    )
+    chatbot = gr.Chatbot(height=550, show_copy_button=True, type="messages")
     
     with gr.Row():
         msg = gr.Textbox(
-            placeholder="Describe the web component or function you need...",
-            show_label=False,
+            placeholder="E.g., Create a glassmorphism login form with Tailwind CSS...", 
+            show_label=False, 
             scale=9
         )
-        submit = gr.Button("Send", variant="primary", scale=1)
-
-    gr.Examples(
-        examples=[
-            "Create a responsive navigation bar using Tailwind CSS.",
-            "Write a JavaScript function to fetch data from an API and display it in a table.",
-            "Design a modern dark mode toggle using CSS variables."
-        ],
-        inputs=msg
-    )
+        submit = gr.Button("Generate", variant="primary", scale=1)
 
     msg.submit(generate_code, [msg, chatbot], [chatbot])
     submit.click(generate_code, [msg, chatbot], [chatbot])