Update app.py

app.py

@@ -1,32 +1,49 @@
+import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 model_name = "Qwen/Qwen2.5-Coder-32B-Instruct"
 
+# Load model and tokenizer (outside the function for efficiency)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype="auto",
-    device_map="auto"
+    device_map="auto",
+    trust_remote_code=True # Add this line for Qwen models
 )
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-prompt = "write a quick sort algorithm."
-messages = [
-    {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},
-    {"role": "user", "content": prompt}
-]
-text = tokenizer.apply_chat_template(
-    messages,
-    tokenize=False,
-    add_generation_prompt=True
-)
-model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
 
-generated_ids = model.generate(
-    **model_inputs,
-    max_new_tokens=512
-)
-generated_ids = [
-    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-]
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) # Add this line for Qwen models
+
+
+@spaces.GPU(required=True)
+def generate_code(prompt):
+    messages = [
+        {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},
+        {"role": "user", "content": prompt}
+    ]
+
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+    generated_ids = model.generate(
+        **model_inputs,
+        max_new_tokens=512
+    )
+
+    generated_ids = [
+        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+    ]
+
+    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return response
 
-response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+# Example usage (optional - remove for Spaces deployment)
+if __name__ == "__main__":
+    prompt = "write a quick sort algorithm."
+    generated_code = generate_code(prompt)
+    print(generated_code)