Spaces:
Running on Zero
Running on Zero
File size: 3,478 Bytes
f507ca1 b242794 f507ca1 83feef8 f507ca1 | 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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | import os
import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL_ID = "Qwen/Qwen2.5-Coder-3B-Instruct"
MAX_INPUT_TOKENS = 8192
MAX_NEW_TOKENS = 512
DEFAULT_SYSTEM_PROMPT = (
"You are a helpful coding assistant. "
"Give accurate, clear answers and use code examples when useful."
)
torch.set_num_threads(min(4, os.cpu_count() or 1))
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
torch_dtype=torch.float16,
attn_implementation="sdpa",
)
model.to("cuda")
model.eval()
def normalize_history(history):
messages = []
for item in history or []:
if isinstance(item, dict):
role = item.get("role")
content = item.get("content")
if role in ("user", "assistant") and isinstance(content, str):
messages.append({"role": role, "content": content})
elif isinstance(item, (list, tuple)) and len(item) == 2:
user_message, assistant_message = item
if isinstance(user_message, str):
messages.append({"role": "user", "content": user_message})
if isinstance(assistant_message, str):
messages.append({"role": "assistant", "content": assistant_message})
return messages
def prepare_inputs(system_prompt, messages):
chat_messages = [{"role": "system", "content": system_prompt}, *messages]
encoded = tokenizer.apply_chat_template(
chat_messages,
tokenize=True,
add_generation_prompt=True,
return_tensors="pt",
return_dict=True,
)
return encoded
@spaces.GPU(size="large", duration=60)
def generate(message, history, system_prompt):
system_prompt = (
DEFAULT_SYSTEM_PROMPT if system_prompt is None else system_prompt
)
messages = normalize_history(history)
messages.append({"role": "user", "content": message})
inputs = prepare_inputs(system_prompt, messages)
while inputs["input_ids"].shape[-1] > MAX_INPUT_TOKENS:
if len(messages) <= 1:
raise gr.Error(
"This message or system prompt is too long. Shorten it and try again."
)
messages = messages[2:] if len(messages) > 2 else messages[1:]
inputs = prepare_inputs(system_prompt, messages)
inputs = {key: value.to("cuda") for key, value in inputs.items()}
with torch.inference_mode():
output = model.generate(
**inputs,
max_new_tokens=MAX_NEW_TOKENS,
do_sample=True,
temperature=0.2,
top_p=0.95,
repetition_penalty=1.05,
pad_token_id=tokenizer.eos_token_id,
eos_token_id=tokenizer.eos_token_id,
)
new_tokens = output[0, inputs["input_ids"].shape[-1]:]
return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
system_prompt_input = gr.Textbox(
label="System prompt",
value=DEFAULT_SYSTEM_PROMPT,
lines=4,
)
demo = gr.ChatInterface(
fn=generate,
chatbot=gr.Chatbot(height=500),
additional_inputs=[system_prompt_input],
title="Qwen2.5 Coder 3B",
description=(
"A coding assistant powered by Qwen2.5-Coder-3B on Hugging Face ZeroGPU."
),
)
demo.queue()
if __name__ == "__main__":
demo.launch()
|