sankar-raul/ICD-10-code-predictor-env

1"""Deterministic graders for the Medical Coding Assistant environment."""2 3from dataclasses import dataclass4 5from .tasks import TaskCase6 7SCORE_EPSILON = 1e-48 9 10@dataclass(frozen=True)11class Submission:12    """Normalized submission shape for scoring."""13 14    primary_code: str15    secondary_codes: tuple[str, ...]16    needs_review: bool17 18 19@dataclass(frozen=True)20class GradeResult:21    """Programmatic grade for a single task."""22 23    task_id: str24    score: float25    feedback: tuple[str, ...]26 27 28def _normalize_code(code: str) -> str:29    return code.strip().upper()30 31 32def _code_family(code: str) -> str:33    normalized = _normalize_code(code)34    if "." in normalized:35        return normalized.split(".", 1)[0]36    if len(normalized) > 3:37        return normalized[:3]38    return normalized39 40 41def _unique_codes(codes: tuple[str, ...]) -> tuple[str, ...]:42    seen: set[str] = set()43    normalized_codes: list[str] = []44    for code in codes:45        normalized = _normalize_code(code)46        if normalized and normalized not in seen:47            seen.add(normalized)48            normalized_codes.append(normalized)49    return tuple(normalized_codes)50 51 52def grade_submission(task: TaskCase, submission: Submission) -> GradeResult:53    """Return a deterministic score strictly in the range (0.0, 1.0)."""54 55    primary = _normalize_code(submission.primary_code)56    secondary = _unique_codes(submission.secondary_codes)57    gold_secondary = _unique_codes(task.gold_secondary)58    feedback: list[str] = []59    score = 0.060 61    if primary == task.gold_primary:62        score += 0.663        feedback.append("Primary diagnosis code is exact.")64    elif primary in task.accepted_primary_alternates:65        score += 0.566        feedback.append("Primary diagnosis is broadly correct but less specific than the gold code.")67    elif _code_family(primary) == _code_family(task.gold_primary):68        score += 0.369        feedback.append("Primary diagnosis is in the right family but not the right code.")70    else:71        feedback.append("Primary diagnosis code is incorrect.")72 73    if gold_secondary:74        matched_secondary = len(set(secondary) & set(gold_secondary))75        secondary_weight = 0.25 / len(gold_secondary)76        if matched_secondary:77            score += matched_secondary * secondary_weight78            feedback.append(79                f"Matched {matched_secondary} supporting code(s) out of {len(gold_secondary)}."80            )81        else:82            feedback.append("Missing the supported secondary code(s).")83 84        extra_secondary = len(set(secondary) - set(gold_secondary))85        if extra_secondary:86            penalty = min(0.15, extra_secondary * 0.05)87            score -= penalty88            feedback.append("Included unsupported secondary code(s).")89 90    if submission.needs_review == task.should_review:91        score += 0.1592        feedback.append("Review/escalation flag is correct.")93    else:94        feedback.append("Review/escalation flag is incorrect.")95 96    score = max(SCORE_EPSILON, min(1.0 - SCORE_EPSILON, round(score, 4)))97    return GradeResult(task_id=task.task_id, score=score, feedback=tuple(feedback))98