Update app.py

app.py

@@ -1,43 +1,36 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+from unsloth import FastModel
 import os
 
-# ✅ Set safe cache directory
 HF_CACHE = "/tmp/hf_cache"
 os.environ["TRANSFORMERS_CACHE"] = HF_CACHE
 os.environ["HF_HOME"] = HF_CACHE
-os.environ["HF_DATASETS_CACHE"] = HF_CACHE
-os.environ["HF_MODULES_CACHE"] = HF_CACHE
 os.makedirs(HF_CACHE, exist_ok=True)
 
 app = FastAPI()
 
 model_id = "srikar-v05/phi3-Mini-Medical-Chat"
-
-tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True, cache_dir=HF_CACHE)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    trust_remote_code=True,
-    cache_dir=HF_CACHE,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+model, tokenizer = FastModel.from_pretrained(
+    model_name = model_id,
+    load_in_4bit = True,
+    max_seq_length = 2048,
 )
 
-# ✅ User input schema
+# Optional: optimise for inference
+FastModel.for_inference(model)
+
 class ChatInput(BaseModel):
     message: str
 
-# ✅ POST endpoint for symptom chat
 @app.post("/chat")
 async def chat_handler(input: ChatInput):
-    # Add provider-like system prompt
-    prompt = f"""You are a kind, attentive oncology provider speaking to a patient. Ask one follow-up question at a time to triage their symptoms.
-
-Patient: {input.message}
-Provider:"""
-
-    inputs = tokenizer(prompt, return_tensors="pt")
+    prompt = (
+        "You are a kind, attentive oncology provider speaking to a patient.\n"
+        "Ask one follow-up question at a time to triage their symptoms.\n\n"
+        f"Patient: {input.message}\nProvider:"
+    )
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     outputs = model.generate(**inputs, max_new_tokens=300)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
     return {"response": response}