add verifier

app.py

@@ -1,5 +1,4 @@
 # app.py
-
 from flask import Flask, render_template, request, send_file, redirect, url_for
 import os
 import logging
@@ -7,7 +6,6 @@ from utils.meldrx import MeldRxAPI
 from utils.oneclick import generate_discharge_paper_one_click
 from huggingface_hub import InferenceClient
 
-
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
@@ -30,7 +28,6 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 if not HF_TOKEN:
     raise ValueError("HF_TOKEN environment variable not set.")
 client = InferenceClient(api_key=HF_TOKEN)
-MODEL_NAME = "meta-llama/Llama-3.3-70B-Instruct"
 
 @app.route('/')
 def index():
@@ -77,7 +74,7 @@ def one_click():
 
         logger.info(f"One-click request - ID: {patient_id}, First: {first_name}, Last: {last_name}, Action: {action}")
 
-        pdf_path, status, basic_summary, ai_summary = generate_discharge_paper_one_click(
+        pdf_path, status, basic_summary, ai_summary, verified_summary = generate_discharge_paper_one_click(
             meldrx_api, client, patient_id, first_name, last_name
         )
 
@@ -86,6 +83,7 @@ def one_click():
                                 status=status,
                                 basic_summary=basic_summary.replace('\n', '<br>') if basic_summary else None,
                                 ai_summary=ai_summary.replace('\n', '<br>') if ai_summary else None,
+                                verified_summary=verified_summary if verified_summary else None,
                                 patient_id=patient_id,
                                 first_name=first_name,
                                 last_name=last_name)
@@ -97,6 +95,7 @@ def one_click():
                             status=status,
                             basic_summary=basic_summary.replace('\n', '<br>') if basic_summary else None,
                             ai_summary=ai_summary.replace('\n', '<br>') if ai_summary else None,
+                            verified_summary=verified_summary if verified_summary else None,
                             patient_id=patient_id,
                             first_name=first_name,
                             last_name=last_name)

requirements.txt

@@ -13,3 +13,4 @@ gradio
 huggingface_hub
 lxml
 reportlab
+lettucedetect

templates/oneclick.html

@@ -2,9 +2,9 @@
 {% block content %}
 <h2>One-Click Discharge Summary</h2>
 <form method="POST">
-    <input type="text" name="patient_id" placeholder="Patient ID (Optional)">
-    <input type="text" name="first_name" placeholder="First Name (Optional)">
-    <input type="text" name="last_name" placeholder="Last Name (Optional)"><br><br>
+    <input type="text" name="patient_id" placeholder="Patient ID (Optional)" value="{{ patient_id or '' }}">
+    <input type="text" name="first_name" placeholder="First Name (Optional)" value="{{ first_name or '' }}">
+    <input type="text" name="last_name" placeholder="Last Name (Optional)" value="{{ last_name or '' }}"><br><br>
     <input type="submit" name="action" value="Display Summary" class="cyberpunk-button">
     <input type="submit" name="action" value="Generate PDF" class="cyberpunk-button">
 </form>
@@ -27,6 +27,13 @@
 </div>
 {% endif %}
 
+{% if verified_summary %}
+<div class="summary-container">
+    <h3>Verified AI Discharge Summary (Hallucinations Highlighted)</h3>
+    <div class="summary-content">{{ verified_summary | safe }}</div>
+</div>
+{% endif %}
+
 <style>
     .status-message {
         margin: 20px 0;

utils/oneclick.py

@@ -3,10 +3,11 @@ from typing import Tuple, Optional, Dict
 from .meldrx import MeldRxAPI
 from .responseparser import PatientDataExtractor
 from .pdfutils import PDFGenerator
+from .verifier import DischargeVerifier  # Import the verifier
 import logging
 import json
 from huggingface_hub import InferenceClient
-import os 
+import os
 
 logger = logging.getLogger(__name__)
 
@@ -15,9 +16,10 @@ if not HF_TOKEN:
     raise ValueError("HF_TOKEN environment variable not set.")
 client = InferenceClient(api_key=HF_TOKEN)
 MODEL_NAME = "meta-llama/Llama-3.3-70B-Instruct"
+verifier = DischargeVerifier()  # Initialize the verifier
 
-def generate_ai_discharge_summary(patient_dict: Dict[str, str], client) -> Optional[str]:
-    """Generate a discharge summary using AI based on extracted patient data."""
+def generate_ai_discharge_summary(patient_dict: Dict[str, str], client) -> Tuple[Optional[str], Optional[str]]:
+    """Generate a discharge summary using AI and verify it for hallucinations."""
     try:
         formatted_summary = format_discharge_summary(patient_dict)
         
@@ -49,12 +51,22 @@ def generate_ai_discharge_summary(patient_dict: Dict[str, str], client) -> Optio
             if content:
                 discharge_summary += content
 
+        discharge_summary = discharge_summary.strip()
         logger.info("AI discharge summary generated successfully")
-        return discharge_summary.strip()
+
+        # Verify the summary for hallucinations
+        question = "Provide a complete discharge summary based on the patient information."
+        verified_summary = verifier.verify_discharge_summary(
+            context=formatted_summary,
+            question=question,
+            answer=discharge_summary
+        )
+        
+        return discharge_summary, verified_summary
 
     except Exception as e:
         logger.error(f"Error generating AI discharge summary: {str(e)}", exc_info=True)
-        return None
+        return None, None
 
 def generate_discharge_paper_one_click(
     api: MeldRxAPI,
@@ -62,12 +74,12 @@ def generate_discharge_paper_one_click(
     patient_id: str = "",
     first_name: str = "",
     last_name: str = ""
-) -> Tuple[Optional[str], str, Optional[str], Optional[str]]:
+) -> Tuple[Optional[str], str, Optional[str], Optional[str], Optional[str]]:
     try:
         patients_data = api.get_patients()
         if not patients_data or "entry" not in patients_data:
             logger.error("No patient data received from MeldRx API")
-            return None, "Failed to fetch patient data from MeldRx API", None, None
+            return None, "Failed to fetch patient data from MeldRx API", None, None, None
 
         logger.debug(f"Raw patient data from API: {patients_data}")
         
@@ -75,7 +87,7 @@ def generate_discharge_paper_one_click(
         
         if not extractor.patients:
             logger.error("No patients found in the parsed data")
-            return None, "No patients found in the data", None, None
+            return None, "No patients found in the data", None, None, None
 
         logger.info(f"Found {len(extractor.patients)} patients in the data")
         
@@ -102,10 +114,8 @@ def generate_discharge_paper_one_click(
             logger.debug(f"Comparing - Input: ID={patient_id_input}, First={first_name_input}, Last={last_name_input}")
             
             matches = True
-            # Only enforce ID match if both input and data have non-empty IDs
             if patient_id_input and patient_id_from_data and patient_id_input != patient_id_from_data:
                 matches = False
-            # Use exact match for names if provided, ignoring case
             if first_name_input and first_name_input != first_name_from_data:
                 matches = False
             if last_name_input and last_name_input != last_name_from_data:
@@ -123,28 +133,28 @@ def generate_discharge_paper_one_click(
             logger.info(f"Available patient names: {all_patient_names}")
             return None, (f"No patients found matching criteria: {search_criteria}\n"
                          f"Available IDs: {', '.join(all_patient_ids)}\n"
-                         f"Available Names: {', '.join(all_patient_names)}"), None, None
-        logger.debug(f"Raw patient data from API: {json.dumps(patients_data, indent=2)}")
+                         f"Available Names: {', '.join(all_patient_names)}"), None, None, None
+        
         patient_data = matching_patients[0]
         logger.info(f"Selected patient data: {patient_data}")
         
         basic_summary = format_discharge_summary(patient_data)
-        ai_summary = generate_ai_discharge_summary(patient_data, client)
+        ai_summary, verified_summary = generate_ai_discharge_summary(patient_data, client)
         
-        if not ai_summary:
-            return None, "Failed to generate AI summary", basic_summary, None
+        if not ai_summary or not verified_summary:
+            return None, "Failed to generate or verify AI summary", basic_summary, None, None
             
         pdf_gen = PDFGenerator()
         filename = f"discharge_{patient_data.get('id', 'unknown')}_{patient_data.get('last_name', 'patient')}.pdf"
         pdf_path = pdf_gen.generate_pdf_from_text(ai_summary, filename)
         
         if pdf_path:
-            return pdf_path, "Discharge summary generated successfully", basic_summary, ai_summary
-        return None, "Failed to generate PDF file", basic_summary, ai_summary
+            return pdf_path, "Discharge summary generated and verified successfully", basic_summary, ai_summary, verified_summary
+        return None, "Failed to generate PDF file", basic_summary, ai_summary, verified_summary
 
     except Exception as e:
         logger.error(f"Error in one-click discharge generation: {str(e)}", exc_info=True)
-        return None, f"Error generating discharge summary: {str(e)}", None, None
+        return None, f"Error generating discharge summary: {str(e)}", None, None, None
 
 def format_discharge_summary(patient_data: dict) -> str:
     """Format patient data into a discharge summary text."""

utils/verifier.py

@@ -0,0 +1,70 @@
+# utils/verifier.py
+from lettucedetect.models.inference import HallucinationDetector
+import logging
+from typing import List, Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+class DischargeVerifier:
+    def __init__(self):
+        """Initialize the hallucination detector."""
+        try:
+            self.detector = HallucinationDetector(
+                method="transformer",
+                model_path="KRLabsOrg/lettucedect-base-modernbert-en-v1",
+            )
+            logger.info("Hallucination detector initialized successfully")
+        except Exception as e:
+            logger.error(f"Failed to initialize hallucination detector: {str(e)}")
+            raise
+
+    def create_interactive_text(self, text: str, spans: List[Dict[str, int | float]]) -> str:
+        """Create interactive HTML with highlighting and hover effects."""
+        html_text = text
+
+        for span in sorted(spans, key=lambda x: x["start"], reverse=True):
+            span_text = text[span["start"]:span["end"]]
+            highlighted_span = (
+                f'<span class="hallucination" title="Confidence: {span["confidence"]:.3f}">{span_text}</span>'
+            )
+            html_text = (
+                html_text[:span["start"]] + highlighted_span + html_text[span["end"]:]
+            )
+
+        return f"""
+        <style>
+            .container {{
+                font-family: Arial, sans-serif;
+                font-size: 16px;
+                line-height: 1.6;
+                padding: 20px;
+            }}
+            .hallucination {{
+                background-color: rgba(255, 99, 71, 0.3);
+                padding: 2px;
+                border-radius: 3px;
+                cursor: help;
+            }}
+            .hallucination:hover {{
+                background-color: rgba(255, 99, 71, 0.5);
+            }}
+        </style>
+        <div class="container">{html_text}</div>
+        """
+
+    def verify_discharge_summary(
+        self, context: str, question: str, answer: str
+    ) -> Optional[str]:
+        """Verify the discharge summary for hallucinations and return highlighted HTML."""
+        try:
+            predictions = self.detector.predict(
+                context=[context],
+                question=question,
+                answer=answer,
+                output_format="spans"
+            )
+            logger.debug(f"Hallucination predictions: {predictions}")
+            return self.create_interactive_text(answer, predictions)
+        except Exception as e:
+            logger.error(f"Error verifying discharge summary: {str(e)}")
+            return None