improve prompt

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from transformers.utils import logging
 
 logging.set_verbosity_debug()
@@ -18,6 +18,7 @@ def load_model():
             model = AutoModelForCausalLM.from_pretrained(
                 "omi-health/sum-small", 
                 trust_remote_code=False,
+                
                 device_map="auto"  # Let the library decide best device mapping
             )
             print(f"GPU: {torch.cuda.get_device_name(0)}")
@@ -40,7 +41,7 @@ def generate_soap_note(doctor_patient_conversation):
     
     try:
         # Create a properly formatted prompt with instructions
-        prompt = f"""<|user|>
+        prompt = f"""<|system|>
 Please generate a structured SOAP (Subjective, Objective, Assessment, Plan) note based on the following doctor-patient conversation:
 
 Include all relevant details in the SOAP note, and ensure that the note is clear and concise. Address each of the following:
@@ -49,8 +50,9 @@ Objective: Observations and findings from the doctor's examination.
 Assessment: Doctor's assessment of the patient's condition.
 Plan: Recommended next steps for the patient's care.
 
-Do not include any additional information or context outside of the SOAP note. Do not include the original prompt or conversation in the output.
-{doctor_patient_conversation}
+Do not include any additional information or context outside of the SOAP note. Do not include the original prompt or conversation in the output.<|end|>
+<|user|>
+{doctor_patient_conversation}<|end|>
 <|assistant|>"""
 
         # Tokenize with reasonable max length
@@ -78,10 +80,8 @@ Do not include any additional information or context outside of the SOAP note. D
         
         # Decode and extract the response part
         decoded_response = tokenizer.batch_decode(generate_ids, skip_special_tokens=True)[0]
-        
-        # Extract only the assistant's response
-        if "<|assistant|>" in decoded_response:
-            decoded_response = decoded_response.split("<|assistant|>")[1].strip()
+
+        logger.debug(decoded_response)
         
         return decoded_response