""" report_generator.py v6 — BioMistral-7B powered ------------------------------------------------ Complete file with all methods intact. """ import logging import re import time from dataclasses import dataclass import torch logger = logging.getLogger(__name__) @dataclass class PatientInfo: age: int gender: str symptoms: str clinical_history: str @dataclass class ImagingInfo: modality: str body_region: str view: str = "" @dataclass class RadiologyReport: examination_details: str = "" clinical_information: str = "" findings: str = "" impression: str = "" recommendations: str = "" confidence_score: float = 0.0 generation_time_seconds: float = 0.0 raw_text: str = "" image_caption: str = "" def to_dict(self) -> dict: return { "examination_details": self.examination_details, "clinical_information": self.clinical_information, "findings": self.findings, "impression": self.impression, "recommendations": self.recommendations, "confidence_score": self.confidence_score, "generation_time_seconds": self.generation_time_seconds, "image_caption": self.image_caption, } def to_text(self) -> str: return "\n".join([ "RADIOLOGY REPORT", "=" * 60, "", "EXAMINATION DETAILS:", self.examination_details, "", "CLINICAL INFORMATION:", self.clinical_information, "", "FINDINGS:", self.findings, "", "IMPRESSION:", self.impression, "", "RECOMMENDATIONS:", self.recommendations, "", f"CONFIDENCE SCORE: {self.confidence_score:.2f}", "", "=" * 60, "END OF REPORT", ]) REGION_FOCUS = { "chest": "lung fields, cardiac silhouette, mediastinum, pleural spaces, bony thorax", "brain": "brain parenchyma, ventricles, sulci, white matter, midline structures, posterior fossa", "abdomen": "liver, spleen, kidneys, pancreas, bowel loops, peritoneal spaces", "pelvis": "pelvic bones, hip joints, bladder, iliac crests, femoral heads, any fixation hardware", "spine": "vertebral bodies, disc spaces, spinal canal, neural foramina, alignment", "knee": "femoral condyles, tibial plateau, patella, joint space, menisci", "shoulder": "glenohumeral joint, rotator cuff, humeral head, acromion", "neck": "cervical spine, soft tissues, thyroid, lymph nodes", "breast": "fibroglandular tissue, lesions, microcalcifications, axillary nodes", "cardiac": "cardiac chambers, great vessels, pericardium", "liver": "hepatic parenchyma, biliary system, portal veins", "thyroid": "thyroid lobes, echogenicity, nodules, vascularity", "obstetric": "fetal anatomy, placenta, amniotic fluid, cervix", } def build_meditron_prompt(patient: PatientInfo, imaging: ImagingInfo, image_caption: str) -> str: view_str = f", {imaging.view} view" if imaging.view else "" focus = REGION_FOCUS.get(imaging.body_region.lower(), f"{imaging.body_region} anatomy and pathology") vision_line = "" if image_caption and len(image_caption) > 20: vision_line = f"\nIMAGE DESCRIPTION: {image_caption}" prompt = ( f"### Instruction:\n" f"You are a consultant radiologist. Write a structured radiology report " f"using standard clinical terminology. Be concise and specific. " f"Only report what is visible. Do not speculate beyond the findings.\n\n" f"STUDY: {imaging.modality} {imaging.body_region}{view_str}\n" f"PATIENT: {patient.age}-year-old {patient.gender}\n" f"CLINICAL INDICATION: {patient.symptoms}\n" f"HISTORY: {patient.clinical_history}\n" f"KEY STRUCTURES: {focus}{vision_line}\n\n" f"Write exactly three sections:\n" f"FINDINGS: (2-4 sentences describing what is seen)\n" f"IMPRESSION: (1-2 sentences — the diagnosis or conclusion)\n" f"RECOMMENDATIONS: (1-2 sentences — next clinical steps)\n\n" f"### Response:\n" f"FINDINGS: The {imaging.body_region} demonstrates" ) return prompt class ReportGenerator: MAX_NEW_TOKENS = 300 TEMPERATURE = 0.1 TOP_P = 0.85 REPETITION_PENALTY = 1.2 def __init__(self, slm_model=None, slm_tokenizer=None): self.model = slm_model self.tokenizer = slm_tokenizer self._device = "cuda" if torch.cuda.is_available() else "cpu" def generate(self, patient: PatientInfo, imaging: ImagingInfo, image_caption: str = "", visual_features=None) -> RadiologyReport: prompt = build_meditron_prompt(patient, imaging, image_caption) start = time.time() raw = self._run_inference(prompt) elapsed = time.time() - start logger.info(f"Raw (250c): {raw[:250]}") report = self._parse(raw, patient, imaging) self._validate_consistency(report, imaging) report.raw_text = raw report.image_caption = image_caption report.generation_time_seconds = round(elapsed, 2) report.confidence_score = self._score(report, imaging) return report # ───────────────────────────────────────────────────────────────── # Inference # ───────────────────────────────────────────────────────────────── def _run_inference(self, prompt: str) -> str: inputs = self.tokenizer( prompt, return_tensors="pt", truncation=True, max_length=800, padding=False, ).to(self._device) with torch.no_grad(): out = self.model.generate( **inputs, max_new_tokens=self.MAX_NEW_TOKENS, temperature=self.TEMPERATURE, top_p=self.TOP_P, repetition_penalty=self.REPETITION_PENALTY, do_sample=True, pad_token_id=self.tokenizer.eos_token_id, eos_token_id=self.tokenizer.eos_token_id, ) new = out[0][inputs["input_ids"].shape[1]:] return self.tokenizer.decode(new, skip_special_tokens=True) # ───────────────────────────────────────────────────────────────── # Text cleaning # ───────────────────────────────────────────────────────────────── def _clean(self, text: str) -> str: text = re.sub(r"```[\s\S]*?```", " ", text) text = re.sub(r"^###.*$", "", text, flags=re.MULTILINE) text = re.sub(r"\n{3,}", "\n\n", text) text = re.sub(r"[ \t]{2,}", " ", text) return text.strip() def _cut_at_sentences(self, text: str, max_sent: int = 4) -> str: sentences = re.split(r'(?<=[.!?])\s+', text.strip()) good = [] for s in sentences: s = s.strip() if len(s) < 8: continue good.append(s) if len(good) >= max_sent: break return " ".join(good) def _extract(self, text: str, section: str, max_sent: int = 3) -> str: aliases = { "findings": ["FINDINGS", "Findings"], "impression": ["IMPRESSION", "Impression", "CONCLUSION", "Conclusion", "DIAGNOSIS"], "recommendations": ["RECOMMENDATIONS", "Recommendations", "RECOMMENDATION", "PLAN", "Plan"], } all_hdrs = [h for v in aliases.values() for h in v] headers = aliases.get(section, [section.upper()]) stopper = "|".join(re.escape(h) for h in all_hdrs if h not in headers) for hdr in headers: pat = (rf"(?:^|\n){re.escape(hdr)}\s*:?\s*(.*?)" rf"(?=\n(?:{stopper})\s*:?|\Z)") m = re.search(pat, text, re.DOTALL | re.IGNORECASE) if m: val = self._clean(m.group(1)) val = self._cut_at_sentences(val, max_sent) val = re.sub(r"\s{2,}", " ", val).strip("- \n") if len(val) > 15: return val return "" # ───────────────────────────────────────────────────────────────── # Parsing # ───────────────────────────────────────────────────────────────── def _parse(self, raw: str, patient: PatientInfo, imaging: ImagingInfo) -> RadiologyReport: report = RadiologyReport() view_str = f", {imaging.view} view" if imaging.view else "" report.examination_details = ( f"{imaging.modality} examination of the {imaging.body_region}" f"{view_str} was performed." ) # Fix double periods (e.g. "infection.." -> "infection.") symptoms_clean = patient.symptoms.rstrip(". ") history_clean = patient.clinical_history.rstrip(". ") report.clinical_information = ( f"{patient.age}-year-old {patient.gender} presenting with " f"{symptoms_clean}. History: {history_clean}." ) full = self._clean( f"FINDINGS: The {imaging.body_region} demonstrates " + raw ) report.findings = self._extract(full, "findings", 3) report.impression = self._extract(full, "impression", 2) report.recommendations = self._extract(full, "recommendations", 2) # Findings fallback if len(report.findings) < 15: lines = [l.strip() for l in full.split("\n") if len(l.strip()) > 30] if lines: candidate = re.sub(r"^FINDINGS\s*:\s*", "", lines[0], flags=re.I) report.findings = self._cut_at_sentences(candidate, 3) # Impression fallback if len(report.impression) < 15: f_lower = report.findings.lower() abnormal_keywords = [ "opacity", "opacities", "ground glass", "consolidation", "lesion", "mass", "nodule", "effusion", "infiltrate", "atelectasis", "fibrosis", "infection", ] normal_keywords = [ "no acute", "no evidence", "unremarkable", "within normal", "normal", "intact", ] if any(w in f_lower for w in abnormal_keywords): report.impression = ( f"Abnormal findings in the {imaging.body_region.lower()} as described. " f"Clinical correlation is recommended." ) elif any(w in f_lower for w in normal_keywords): report.impression = ( f"No acute {imaging.body_region.lower()} abnormality identified. " f"Clinical correlation is recommended." ) else: report.impression = ( f"Findings in the {imaging.body_region.lower()} as described. " f"Clinical correlation is recommended." ) # Recommendations fallback if len(report.recommendations) < 15: report.recommendations = ( "Clinical correlation with the patient's history and examination " "findings is recommended. Follow-up imaging as clinically indicated." ) return report # ───────────────────────────────────────────────────────────────── # Consistency check # ───────────────────────────────────────────────────────────────── def _validate_consistency(self, report: RadiologyReport, imaging: ImagingInfo): f = report.findings.lower() imp = report.impression.lower() abnormal_keywords = [ "opacity", "opacities", "ground glass", "consolidation", "lesion", "mass", "nodule", "effusion", "infiltrate", ] if any(w in f for w in abnormal_keywords): if any(x in imp for x in ["no acute", "normal", "unremarkable"]): report.impression = self._build_specific_impression( report, imaging ) def _build_specific_impression(self, report: RadiologyReport, imaging: ImagingInfo) -> str: """ Builds a specific, clinically meaningful impression by: 1. Extracting the key finding from FINDINGS 2. Incorporating clinical history context 3. Suggesting a differential diagnosis """ f = report.findings.lower() hist = report.clinical_information.lower() region = imaging.body_region.lower() # ── Identify key finding ─────────────────────────────────── if "ground glass" in f or "ggo" in f: key_finding = "Ground-glass opacities" if "bilateral" in f or "both lung" in f: key_finding = "Bilateral ground-glass opacities" elif "consolidation" in f: key_finding = "Airspace consolidation" if "bilateral" in f: key_finding = "Bilateral consolidation" elif "right" in f: key_finding = "Right-sided consolidation" elif "left" in f: key_finding = "Left-sided consolidation" elif "opacity" in f or "opacities" in f: key_finding = "Pulmonary opacities" if "bilateral" in f: key_finding = "Bilateral pulmonary opacities" elif "nodule" in f: key_finding = "Pulmonary nodule(s)" elif "effusion" in f: key_finding = "Pleural effusion" if "bilateral" in f: key_finding = "Bilateral pleural effusions" elif "mass" in f or "lesion" in f: key_finding = f"Focal {region} lesion" elif "infiltrate" in f: key_finding = "Pulmonary infiltrates" else: key_finding = f"Abnormal {region} findings" # ── Clinical context from history ────────────────────────── context = "" if any(w in hist for w in ["immunosuppress", "immunocompromis", "mycophenol", "steroid", "transplant", "hiv", "aids", "chemotherapy"]): context = ( f" In the context of immunosuppression, opportunistic " f"infection (e.g. PCP, fungal, CMV pneumonitis) should " f"be considered." ) elif any(w in hist for w in ["infection", "pneumonia", "fever", "cough"]): context = ( f" Findings are consistent with an infective or " f"inflammatory aetiology." ) elif any(w in hist for w in ["malignancy", "cancer", "tumour", "tumor", "oncology", "lymphoma"]): context = ( f" In the context of known malignancy, metastatic " f"disease should be considered." ) elif any(w in hist for w in ["autoimmune", "lupus", "rheumatoid", "sarcoid", "vasculitis"]): context = ( f" In the context of autoimmune disease, an inflammatory " f"or drug-related aetiology should be considered." ) # ── Build impression ─────────────────────────────────────── impression = ( f"{key_finding} as described above, likely representing " f"an infective or inflammatory process.{context} " f"Clinical correlation is recommended." ) return impression # ───────────────────────────────────────────────────────────────── # Confidence scoring # ───────────────────────────────────────────────────────────────── def _score(self, report: RadiologyReport, imaging: ImagingInfo) -> float: score = 0.45 f = report.findings if len(f) > 40: score += 0.10 if len(f) > 100: score += 0.05 anatomy = { "brain": ["parenchyma", "ventricle", "sulci", "cortex", "signal", "midline", "fossa", "white matter"], "chest": ["lung", "cardiac", "pleural", "mediastinum", "consolidation", "silhouette", "costophrenic"], "pelvis": ["femoral", "acetabul", "iliac", "hip", "sacrum", "hardware", "screw", "fixation", "prosthes"], "abdomen": ["liver", "spleen", "kidney", "bowel", "pancreas", "peritoneal", "aorta"], "spine": ["vertebr", "disc", "canal", "foramina", "alignment", "cord", "facet"], "knee": ["condyle", "meniscus", "patella", "tibial", "ligament", "joint space"], } region_key = imaging.body_region.lower() terms = anatomy.get(region_key, []) matched = sum(1 for t in terms if t in f.lower()) score += min(matched * 0.03, 0.15) imp = report.impression if len(imp) > 30: score += 0.08 good = ["consistent with", "no evidence", "no acute", "identified", "findings", "normal", "unremarkable", "in keeping with"] if any(p in imp.lower() for p in good): score += 0.08 if "clinical presentation of a" not in imp: score += 0.05 if len(report.recommendations) > 20: score += 0.04 return round(max(0.45, min(0.88, score)), 2)