|
|
|
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 |
