Update app.py

app.py

@@ -7,7 +7,7 @@ from transformers import (
     AutoModelForCausalLM
 )
 
-MODEL_NAME = "Qwen/Qwen2.5-3B-Instruct"
+MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
 
 print("Loading model...")
 
@@ -18,9 +18,9 @@ tokenizer = AutoTokenizer.from_pretrained(
 
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
-    torch_dtype="auto",
-    device_map="auto",
-    trust_remote_code=True
+    torch_dtype=torch.float32,
+    trust_remote_code=True,
+    low_cpu_mem_usage=True
 )
 
 print("Model loaded successfully.")
@@ -132,7 +132,11 @@ def generate_response(prompt):
         max_length=4096
     )
 
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    model = model.to(device)
+    
+    inputs = {k: v.to(device) for k, v in inputs.items()}
 
     outputs = model.generate(
         **inputs,