Upload folder using huggingface_hub

README.md

@@ -9,7 +9,7 @@ app_file: app.py
 pinned: false
 license: apache-2.0
 short_description: Code assistant powered by fine-tuned Qwen 2.5 Coder
-suggested_hardware: t4-small
+suggested_hardware: zero-a10g
 ---
 
 # Qwen 2.5 Coder Assistant

app.py

@@ -1,19 +1,21 @@
 import gradio as gr
+import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 MODEL_ID = "GhostScientist/qwen25-coder-1.5b-codealpaca-sft"
 
-# Load model and tokenizer
+# Load tokenizer at startup
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+# Load model at startup (will be moved to GPU when @spaces.GPU is called)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.float16,
-    device_map="auto",
 )
 
-def respond(message, history, system_message, max_tokens, temperature, top_p):
+@spaces.GPU
+def generate_response(message, history, system_message, max_tokens, temperature, top_p):
     """Generate response using the fine-tuned Qwen coder model."""
     messages = [{"role": "system", "content": system_message}]
 
@@ -33,12 +35,9 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
     )
     inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
-    # Set up streaming
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-
-    generation_kwargs = dict(
+    # Generate response
+    outputs = model.generate(
         **inputs,
-        streamer=streamer,
         max_new_tokens=int(max_tokens),
         temperature=temperature,
         top_p=top_p,
@@ -46,15 +45,9 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
         pad_token_id=tokenizer.eos_token_id,
     )
 
-    # Run generation in a thread
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-
-    # Stream the response
-    response = ""
-    for new_text in streamer:
-        response += new_text
-        yield response
+    # Decode only the new tokens
+    response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+    return response
 
 
 SYSTEM_PROMPT = """You are an expert coding assistant. You help users write, debug, explain, and improve code.
@@ -70,7 +63,7 @@ EXAMPLES = [
 ]
 
 demo = gr.ChatInterface(
-    respond,
+    generate_response,
     title="Qwen 2.5 Coder Assistant",
     description="""A fine-tuned Qwen 2.5 Coder 1.5B model for code assistance.
     Ask me to write code, explain concepts, debug issues, or help with any programming task!