Update app.py

Edit the response not to include prompt

app.py

@@ -40,21 +40,19 @@ def get_clinical_code(clinical_text):
     """
     if not model_loaded:
         return "Error: The model could not be loaded. Please check the logs."
-
     if not clinical_text:
         return "Please enter some clinical text."
 
     # This is our prompt template. It's designed to guide the model
     # to perform the specific task of clinical coding.
     # We are asking for an ICD-10 code, which is a common standard.
-    prompt = f"""
-    <start_of_turn>user
-    You are an expert medical coder. Your task is to analyze the following clinical text and determine the most appropriate ICD-10 code. Provide only the ICD-10 code and a brief description.
-    Clinical Text: "{clinical_text}"
-    Provide the ICD-10 code and a brief description.
-    <end_of_turn>
-    <start_of_turn>model
-    """
+    prompt = f"""<start_of_turn>user
+You are an expert medical coder. Your task is to analyze the following clinical text and determine the most appropriate ICD-10 code. Provide only the ICD-10 code and a brief description.
+Clinical Text: "{clinical_text}"
+Provide the ICD-10 code and a brief description.
+<end_of_turn>
+<start_of_turn>model
+"""
     # Prepare the input for the model
     # It's good practice to ensure the input_ids are on the correct device.
     # model.device will give you the device where the model currently resides (GPU if available).
@@ -70,15 +68,16 @@ def get_clinical_code(clinical_text):
         temperature=0.7, # A lower temperature makes the output more deterministic
     )
 
-    # Decode the output and clean it up
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Extract the relevant part of the response
-    # The model will output the prompt as well, so we need to remove it.
-    model_response_start = response.find("<start_of_turn>model") + len("<start_of_turn>model")
-    clean_response = response[model_response_start:].strip()
+    # Decode only the newly generated tokens
+    # We slice the outputs tensor to exclude the input_ids (the prompt)
+    generated_tokens = outputs[0, input_ids.input_ids.shape[1]:]
+    response = tokenizer.decode(generated_tokens, skip_special_tokens=True)
 
-    return clean_response
+    # The prompt explicitly asks for "only the ICD-10 code and a brief description."
+    # So, we expect the model to start directly with the code and description.
+    # No further slicing based on "<start_of_turn>model" should be needed for *this* part.
+    # We can just return the decoded response.
+    return response.strip()
 
 # Create the Gradio Interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
@@ -89,7 +88,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         *Disclaimer: This is a demonstration and not for professional medical use.*
         """
     )
-
     with gr.Row():
         # Input Textbox
         input_text = gr.Textbox(
@@ -97,7 +95,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             placeholder="e.g., Patient presents with a severe headache and photophobia...",
             lines=10
         )
-
         # Output Textbox
         output_text = gr.Textbox(
             label="Suggested Clinical Code (ICD-10)",
@@ -112,7 +109,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         inputs=input_text,
         outputs=output_text
     )
-
     gr.Examples(
         examples=[
             ["The patient complains of a persistent cough and fever for the past three days. Chest X-ray shows signs of pneumonia."],