Spaces:
Runtime error
Runtime error
| """ | |
| 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__) | |
| class PatientInfo: | |
| age: int | |
| gender: str | |
| symptoms: str | |
| clinical_history: str | |
| class ImagingInfo: | |
| modality: str | |
| body_region: str | |
| view: str = "" | |
| 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) | |