Update app.py

app.py

@@ -4,34 +4,40 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import os
 
-# Set Hugging Face cache to a safe location in Hugging Face Space
-os.environ["HF_HOME"] = "/tmp"
-os.environ["HF_DATASETS_CACHE"] = "/tmp"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp"
-os.environ["HF_MODULES_CACHE"] = "/tmp"
-
-
+# ✅ 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()
 
-# Load ChatDoctor model
 model_id = "Dashanka/medical-chatbot-Llama3.1-8B-instruct-4bit"
-tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+
+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
 )
 
-
-
-
+# ✅ User input schema
 class ChatInput(BaseModel):
     message: str
 
+# ✅ POST endpoint for symptom chat
 @app.post("/chat")
 async def chat_handler(input: ChatInput):
-    inputs = tokenizer(input.message, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=200)
+    # 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")
+    outputs = model.generate(**inputs, max_new_tokens=300)
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return {"response": response}