Spaces:
Running on Zero
Running on Zero
| import logging | |
| import threading | |
| import gradio as gr | |
| import spaces | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer | |
| logging.basicConfig(level=logging.INFO) | |
| MODEL_ID = "Qwen/Qwen2.5-7B-Instruct" | |
| SYSTEM_PROMPT = ( | |
| "You are Tony-Flash, a helpful AI assistant. " | |
| "Answer in the user's language. Be accurate and clear." | |
| ) | |
| # Load at startup; ZeroGPU attaches the real GPU only inside @spaces.GPU calls. | |
| tok = AutoTokenizer.from_pretrained(MODEL_ID) | |
| model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16) | |
| model.to("cuda") | |
| def chat(message, history): | |
| if not message or not message.strip(): | |
| yield "Please enter a message." | |
| return | |
| messages = [{"role": "system", "content": SYSTEM_PROMPT}] | |
| for item in history or []: | |
| if isinstance(item, dict) and item.get("role") in ("user", "assistant"): | |
| content = item.get("content") | |
| if isinstance(content, str): | |
| messages.append({"role": item["role"], "content": content}) | |
| messages.append({"role": "user", "content": message.strip()}) | |
| inputs = tok.apply_chat_template( | |
| messages, | |
| add_generation_prompt=True, | |
| return_tensors="pt", | |
| return_dict=True, | |
| ).to("cuda") | |
| streamer = TextIteratorStreamer(tok, skip_prompt=True, skip_special_tokens=True) | |
| gen_kwargs = dict( | |
| **inputs, | |
| streamer=streamer, | |
| max_new_tokens=1024, | |
| do_sample=True, | |
| temperature=0.7, | |
| ) | |
| threading.Thread(target=model.generate, kwargs=gen_kwargs).start() | |
| answer = "" | |
| for piece in streamer: | |
| answer += piece | |
| yield answer | |
| demo = gr.ChatInterface( | |
| fn=chat, | |
| title="🧠 Tony-Flash", | |
| description=f"Chat with {MODEL_ID} on ZeroGPU", | |
| textbox=gr.Textbox( | |
| placeholder="اكتب رسالتك هنا... / Type your message", | |
| container=False, | |
| ), | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue().launch(show_error=True, ssr_mode=False) |