Spaces:
Running on Zero
Running on Zero
| 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 | |
| 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(type="messages", 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() | |