Update app.py

app.py

@@ -12,18 +12,18 @@ ADAPTER_MODEL_ID = "vsple/LegalBuddy-Qwen-1.5B"
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, trust_remote_code=True)
 
-# Load base model
-print("Loading base model...")
+# Load model and adapter
+print("Loading model and adapter...")
+device = "cuda" if torch.cuda.is_available() else "cpu"
 base_model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL_ID,
     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto",
+    low_cpu_mem_usage=True,
     trust_remote_code=True
 )
 
-# Load adapter
-print("Loading adapter...")
 model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL_ID)
+model = model.to(device)
 model = model.eval()
 
 def predict(message, history, system_prompt, max_tokens, temperature, top_p):