Add prompt

app.py

@@ -13,21 +13,37 @@ def load_model():
     model = AutoModelForCausalLM.from_pretrained("omi-health/sum-small", trust_remote_code=False)
     
     return model, tokenizer
-    return model, tokenizer
 
 # Function to generate SOAP notes
 def generate_soap_note(doctor_patient_conversation):
     if not doctor_patient_conversation.strip():
         return "Please enter a doctor-patient conversation."
-        
-    # Tokenize and generate
-    inputs = tokenizer(doctor_patient_conversation, return_tensors="pt")
     
-    generate_ids = model.generate(inputs.input_ids, max_length=200, num_beams=5, no_repeat_ngram_size=2, early_stopping=True)
+    # Create a properly formatted prompt with instructions
+    prompt = f"""<|user|>
+Please generate a structured SOAP (Subjective, Objective, Assessment, Plan) note based on the following doctor-patient conversation:
 
+{doctor_patient_conversation}
+<|assistant|>"""
+    
+    # Tokenize and generate
+    inputs = tokenizer(prompt, return_tensors="pt")
+    
+    generate_ids = model.generate(
+        inputs.input_ids, 
+        max_length=2048,
+        num_beams=5, 
+        no_repeat_ngram_size=2, 
+        early_stopping=True
+    )
 
     # Decode and extract the response part
     decoded_response = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    
+    # Extract only the assistant's response (remove the prompt part)
+    if "<|assistant|>" in decoded_response:
+        decoded_response = decoded_response.split("<|assistant|>")[1].strip()
+    
     logger.debug(f"Decoded response: {decoded_response}")
     return decoded_response