Spaces:
Running on Zero
Running on Zero
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import json | |
| import spaces | |
| import torch | |
| import gradio as gr | |
| from threading import Thread | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer | |
| BUILD = "tools-v1" | |
| MODEL_ID = "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF" | |
| GGUF_FILE = "qwen2.5-coder-7b-instruct-q4_k_m.gguf" | |
| SYSTEM_PROMPT = "You are Qwen2.5-Coder, a helpful and concise coding assistant." | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, gguf_file=GGUF_FILE) | |
| model = ( | |
| AutoModelForCausalLM.from_pretrained( | |
| MODEL_ID, | |
| gguf_file=GGUF_FILE, | |
| dtype=torch.float32, | |
| attn_implementation="sdpa", | |
| ) | |
| .to("cuda") | |
| .eval() | |
| ) | |
| def _build_prompt(message, history): | |
| messages = [{"role": "system", "content": SYSTEM_PROMPT}] | |
| for turn in history: | |
| role, content = turn.get("role"), turn.get("content") | |
| if role in ("user", "assistant") and isinstance(content, str) and content: | |
| messages.append({"role": role, "content": content}) | |
| messages.append({"role": "user", "content": message}) | |
| return tokenizer.apply_chat_template( | |
| messages, tokenize=False, add_generation_prompt=True | |
| ) | |
| def _stream(message, history, max_new_tokens, temperature, top_p): | |
| text = _build_prompt(message, history) | |
| inputs = tokenizer(text, return_tensors="pt").to("cuda") | |
| streamer = TextIteratorStreamer( | |
| tokenizer, skip_prompt=True, skip_special_tokens=True | |
| ) | |
| kwargs = dict( | |
| **inputs, | |
| max_new_tokens=int(max_new_tokens), | |
| do_sample=temperature > 0, | |
| temperature=max(float(temperature), 1e-5), | |
| top_p=float(top_p), | |
| streamer=streamer, | |
| pad_token_id=tokenizer.eos_token_id, | |
| ) | |
| Thread(target=model.generate, kwargs=kwargs).start() | |
| partial = "" | |
| for chunk in streamer: | |
| partial += chunk | |
| yield partial | |
| def respond(chat_history, temperature, top_p, max_new_tokens): | |
| message = chat_history[-1]["content"] | |
| history = chat_history[:-1] | |
| chat_history = chat_history + [{"role": "assistant", "content": ""}] | |
| for partial in _stream(message, history, max_new_tokens, temperature, top_p): | |
| chat_history[-1]["content"] = partial | |
| yield chat_history | |
| def add_user(message, chat_history): | |
| if not message or not message.strip(): | |
| return "", chat_history | |
| return "", chat_history + [{"role": "user", "content": message}] | |
| def _to_qwen_messages(messages): | |
| out = [] | |
| for m in messages or []: | |
| role = m.get("role") | |
| content = m.get("content") or "" | |
| if role == "assistant" and m.get("tool_calls"): | |
| parts = [content] if content else [] | |
| for tc in m["tool_calls"]: | |
| fn = (tc or {}).get("function", {}) or {} | |
| args = fn.get("arguments") | |
| if isinstance(args, str): | |
| try: | |
| args = json.loads(args) | |
| except Exception: | |
| args = {} | |
| parts.append( | |
| "<tool_call>\n" | |
| + json.dumps( | |
| {"name": fn.get("name"), "arguments": args}, | |
| ensure_ascii=False, | |
| ) | |
| + "\n</tool_call>" | |
| ) | |
| content = "\n".join(parts) | |
| if role in ("system", "user", "assistant", "tool"): | |
| out.append({"role": role, "content": content}) | |
| return out | |
| def tool_chat(messages_json, tools_json, temperature, top_p, max_new_tokens): | |
| messages = json.loads(messages_json) if messages_json else [] | |
| tools = json.loads(tools_json) if tools_json else None | |
| chat = _to_qwen_messages(messages) | |
| if not chat or chat[0]["role"] != "system": | |
| chat = [{"role": "system", "content": SYSTEM_PROMPT}] + chat | |
| kwargs = {} | |
| if tools: | |
| kwargs["tools"] = tools | |
| try: | |
| text = tokenizer.apply_chat_template( | |
| chat, tokenize=False, add_generation_prompt=True, **kwargs | |
| ) | |
| except TypeError: | |
| text = tokenizer.apply_chat_template( | |
| chat, tokenize=False, add_generation_prompt=True | |
| ) | |
| inputs = tokenizer(text, return_tensors="pt").to("cuda") | |
| out = model.generate( | |
| **inputs, | |
| max_new_tokens=int(max_new_tokens), | |
| do_sample=float(temperature) > 0, | |
| temperature=max(float(temperature), 1e-5), | |
| top_p=float(top_p), | |
| pad_token_id=tokenizer.eos_token_id, | |
| ) | |
| return tokenizer.decode( | |
| out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True | |
| ) | |
| with gr.Blocks(title="Qwen2.5 Coder 7B (GGUF)") as demo: | |
| gr.Markdown( | |
| "# Qwen2.5-Coder-7B-Instruct (GGUF)\n" | |
| "Streaming chat demo running on **ZeroGPU**. " + BUILD | |
| ) | |
| chatbot = gr.Chatbot( | |
| type="messages", height=440, label="Qwen2.5-Coder", allow_tags=False | |
| ) | |
| with gr.Row(): | |
| msg = gr.Textbox( | |
| placeholder="Ask me to write or explain code...", | |
| scale=8, | |
| container=False, | |
| ) | |
| send = gr.Button("Send", variant="primary", scale=1) | |
| with gr.Row(): | |
| temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="Temperature") | |
| top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="Top-p") | |
| max_tokens = gr.Slider(64, 2048, value=512, step=64, label="Max new tokens") | |
| clear = gr.Button("Clear") | |
| gr.Examples( | |
| [ | |
| ["Write a Python function to merge two sorted lists."], | |
| ["Debug this SQL: SELECT * FORM users;"], | |
| ], | |
| inputs=msg, | |
| ) | |
| with gr.Row(visible=False): | |
| tc_msgs = gr.Textbox() | |
| tc_tools = gr.Textbox() | |
| tc_temp = gr.Number(value=0.2) | |
| tc_top_p = gr.Number(value=0.9) | |
| tc_max = gr.Number(value=512) | |
| tc_btn = gr.Button("tool") | |
| tc_out = gr.Textbox() | |
| tc_btn.click( | |
| tool_chat, | |
| [tc_msgs, tc_tools, tc_temp, tc_top_p, tc_max], | |
| tc_out, | |
| api_name="tool_chat", | |
| ) | |
| msg.submit(add_user, [msg, chatbot], [msg, chatbot], queue=False).then( | |
| respond, [chatbot, temperature, top_p, max_tokens], chatbot | |
| ) | |
| send.click(add_user, [msg, chatbot], [msg, chatbot], queue=False).then( | |
| respond, [chatbot, temperature, top_p, max_tokens], chatbot | |
| ) | |
| clear.click(lambda: None, None, chatbot, queue=False) | |
| demo.queue() | |
| demo.launch(mcp_server=True) | |