Update app.py

app.py

@@ -12,18 +12,16 @@ 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...")
-base_model = AutoModelForCausalLM.from_pretrained(
+# Load model (Fixed loading logic for HF Spaces)
+print("Loading model and adapter...")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL_ID,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto",
+    torch_dtype=torch.float32,
     trust_remote_code=True
-)
+).to(device)
 
-# Load adapter
-print("Loading adapter...")
-model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL_ID)
+model = PeftModel.from_pretrained(model, ADAPTER_MODEL_ID)
 model = model.eval()
 
 def predict(message, history, system_prompt, max_tokens, temperature, top_p):