Spaces:
Runtime error
Runtime error
File size: 3,220 Bytes
c0bd465 2fabed1 c0bd465 2fabed1 c0bd465 c96232a 2fabed1 148451f c96232a 148451f 9fafb18 2fabed1 148451f 2fabed1 148451f 2fabed1 148451f 2fabed1 c96232a 9fafb18 2fabed1 c96232a 2fabed1 c96232a 9fafb18 2fabed1 c96232a 148451f 2fabed1 148451f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | import os
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
hf_token = os.getenv("HF_TOKEN")
model_id = "ZyperAI/Z-AI-0.1-1.1B-Code.web"
print("Loading model and tokenizer...")
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
)
print("Model loaded successfully.")
def generate_code(prompt, history):
# 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})
# 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(
input_ids=inputs, # Fix 3: Transformers generation kwargs expects 'input_ids', not 'inputs'
streamer=streamer,
max_new_tokens=1024,
do_sample=True,
temperature=0.7,
top_p=0.9
)
thread = Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()
# 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:
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**")
chatbot = gr.Chatbot(height=500, show_copy_button=True, type="messages")
with gr.Row():
msg = gr.Textbox(
placeholder="E.g., Create a responsive navigation bar with CSS...",
show_label=False,
scale=9
)
submit = gr.Button("Build", variant="primary", scale=1)
# 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()
|