Update model_utils.py

model_utils.py

@@ -14,10 +14,16 @@ from loader import load_curriculum, load_manual_qa, rebuild_combined_qa
 MODEL_NAME = "SeaLLMs/SeaLLMs-v3-1.5B-Chat"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
     torch_dtype=torch.float32,
-)
+).to(device)
+
+model.eval()
+
 
 # Load data once at import time
 load_curriculum()
@@ -109,7 +115,7 @@ def build_prompt(question: str) -> str:
 
 def generate_answer(question: str) -> str:
     prompt = build_prompt(question)
-    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     with torch.no_grad():
         outputs = model.generate(
             **inputs,