Upload app.py

app.py

@@ -3,7 +3,6 @@ from pydantic import BaseModel
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from peft import PeftModel
-
 import os
 from huggingface_hub import login
 
@@ -14,35 +13,44 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
     login(token=HF_TOKEN)
 
-
 # ✅ Initialize FastAPI
 app = FastAPI()
 
-# ✅ Define Base Model & LoRA Adapter Repository
-base_model_name = "mistralai/Mistral-7B-v0.1"
+# ✅ Define Base Model & LoRA Adapter Repository (Smaller Model for Hugging Face Spaces)
+base_model_name = "mistralai/Mistral-7B-Instruct-v0.2"  # 🔹 Using a smaller model
 lora_repo_id = "khushi1234455687/fine-tuned-medical-qa-V8"
 
+# ✅ Automatically Select CPU (Hugging Face Spaces Does NOT Support GPU)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 # ✅ Load Tokenizer
 tokenizer = AutoTokenizer.from_pretrained(base_model_name)
 
 # ✅ Configure 4-bit Quantization
 quantization_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    llm_int8_enable_fp32_cpu_offload=True,
-    offload_buffers=True
+    load_in_4bit=True
 )
 
-# ✅ Load Base Model
-base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_name,
-    quantization_config=quantization_config,
-    device_map="auto",
-    torch_dtype=torch.float16
-)
+# ✅ Load Base Model (Optimized for CPU)
+try:
+    base_model = AutoModelForCausalLM.from_pretrained(
+        base_model_name,
+        quantization_config=quantization_config,
+        device_map="auto",  # ✅ Automatically assigns layers to CPU
+        torch_dtype=torch.float16
+    )
+except Exception as e:
+    print(f"❌ Error loading base model: {e}")
+    raise e
 
 # ✅ Load LoRA Adapter
-model = PeftModel.from_pretrained(base_model, lora_repo_id)
-model.eval()
+try:
+    model = PeftModel.from_pretrained(base_model, lora_repo_id)
+    model.to(device)  # ✅ Ensure model is on the correct device
+    model.eval()
+except Exception as e:
+    print(f"❌ Error loading LoRA adapter: {e}")
+    raise e
 
 print("✅ Model is loaded and API is ready!")
 
@@ -53,10 +61,17 @@ class QueryRequest(BaseModel):
 @app.post("/generate")
 async def generate_answer(request: QueryRequest):
     """Generate an answer for a given medical question."""
-    inputs = tokenizer(request.question, return_tensors="pt").to("cuda")
-    with torch.no_grad():
-        output = model.generate(**inputs, max_length=256)
-    answer = tokenizer.decode(output[0], skip_special_tokens=True)
-    
+    try:
+        inputs = tokenizer(request.question, return_tensors="pt").to(device)  # ✅ Move to device
+        with torch.no_grad():
+            output = model.generate(**inputs, max_length=256)
+        answer = tokenizer.decode(output[0], skip_special_tokens=True)
+    except Exception as e:
+        return {"error": str(e)}
+
     return {"question": request.question, "answer": answer}
 
+# ✅ Health Check Endpoint
+@app.get("/health")
+async def health_check():
+    return {"status": "running", "device": device}