...

llm_engine.py

@@ -4,6 +4,7 @@ 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,
@@ -11,8 +12,10 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto"
 )
 
+# Create generation pipeline
 generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
+# Handler functions
 def handle_symptoms(symptoms: str) -> str:
     prompt = SYMPTOM_PROMPT_TEMPLATE.format(symptoms=symptoms)
     output = generator(prompt, max_new_tokens=512, do_sample=True)[0]["generated_text"]