...

llm_engine.py

@@ -4,24 +4,29 @@ from prompts import SYMPTOM_PROMPT_TEMPLATE, QUESTION_PROMPT_TEMPLATE
 
 MODEL_NAME = "GEMINI-Lab/MedicalGPT-LLAMA2-7B"
 
-# Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto"
-)
+# Load model & tokenizer
+try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        device_map="auto"
+    )
+    generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+except Exception as e:
+    print(f"Model loading failed: {e}")
+    generator = None
 
-# Create generation pipeline
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-# Handler functions
 def handle_symptoms(symptoms: str) -> str:
+    if generator is None:
+        return "Model failed to load."
     prompt = SYMPTOM_PROMPT_TEMPLATE.format(symptoms=symptoms)
     output = generator(prompt, max_new_tokens=512, do_sample=True)[0]["generated_text"]
     return output[len(prompt):].strip()
 
 def handle_question(question: str) -> str:
+    if generator is None:
+        return "Model failed to load."
     prompt = QUESTION_PROMPT_TEMPLATE.format(question=question)
     output = generator(prompt, max_new_tokens=512, do_sample=True)[0]["generated_text"]
     return output[len(prompt):].strip()