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

1"""Build a searchable ICD codebook model from ICDCodeSet.csv."""2 3from __future__ import annotations4 5import argparse6import csv7import math8import pickle9import re10from collections import Counter, defaultdict11from dataclasses import dataclass12from pathlib import Path13 14 15TOKEN_RE = re.compile(r"[a-z0-9]+")16STOPWORDS = {17    "the",18    "and",19    "of",20    "to",21    "in",22    "for",23    "with",24    "due",25    "other",26    "unspecified",27}28 29 30def _tokenize(text: str) -> list[str]:31    return TOKEN_RE.findall(text.lower())32 33 34@dataclass(frozen=True)35class ICDEntry:36    code: str37    description: str38 39 40class ICDCodebookModel:41    """Simple lexical retrieval model for ICD descriptions."""42 43    def __init__(self) -> None:44        self.entries: list[ICDEntry] = []45        self.code_to_desc: dict[str, str] = {}46        self.token_to_codes: dict[str, Counter[str]] = defaultdict(Counter)47        self.code_token_counts: dict[str, Counter[str]] = {}48        self.doc_freq: Counter[str] = Counter()49        self.num_docs = 050 51    def fit(self, entries: list[ICDEntry]) -> None:52        self.entries = entries53        self.code_to_desc = {e.code: e.description for e in entries}54        self.num_docs = len(entries)55        for entry in entries:56            tokens = [t for t in _tokenize(entry.description) if t not in STOPWORDS]57            token_counts = Counter(tokens)58            self.code_token_counts[entry.code] = token_counts59            for token in token_counts:60                self.doc_freq[token] += 161            for token, count in token_counts.items():62                self.token_to_codes[token][entry.code] += count63 64    def _idf(self, token: str) -> float:65        # Smoothed IDF.66        return math.log((1 + self.num_docs) / (1 + self.doc_freq[token])) + 1.067 68    def search(self, text: str, top_k: int = 5) -> list[tuple[str, float]]:69        query_tokens = [t for t in _tokenize(text) if t not in STOPWORDS]70        if not query_tokens:71            return []72        query_counts = Counter(query_tokens)73        scores: Counter[str] = Counter()74        for token, q_count in query_counts.items():75            idf = self._idf(token)76            for code, tf in self.token_to_codes.get(token, {}).items():77                scores[code] += (q_count * idf) * (tf * idf)78        if not scores:79            return []80        ranked = sorted(scores.items(), key=lambda x: (-x[1], x[0]))81        return [(code, float(score)) for code, score in ranked[:top_k]]82 83 84def load_codeset(csv_path: Path) -> list[ICDEntry]:85    entries: list[ICDEntry] = []86    with csv_path.open("r", encoding="utf-8", newline="") as f:87        reader = csv.DictReader(f)88        for row in reader:89            code = row["ICDCode"].strip().upper()90            desc = row["Description"].strip()91            if code and desc:92                entries.append(ICDEntry(code=code, description=desc))93    return entries94 95 96def save_model(model: ICDCodebookModel, pkl_path: Path) -> None:97    pkl_path.parent.mkdir(parents=True, exist_ok=True)98    with pkl_path.open("wb") as f:99        pickle.dump(model, f)100 101 102def main() -> None:103    parser = argparse.ArgumentParser()104    parser.add_argument("--csv", default="data/ICDCodeSet.csv", help="Path to ICDCodeSet.csv")105    parser.add_argument(106        "--save-pkl",107        default="outputs/icd_codeset_model.pkl",108        help="Where to save the trained ICD codebook model.",109    )110    parser.add_argument(111        "--sample-query",112        default="cholera due to vibrio",113        help="Optional query to test retrieval output.",114    )115    args = parser.parse_args()116 117    entries = load_codeset(Path(args.csv))118    model = ICDCodebookModel()119    model.fit(entries)120    save_model(model, Path(args.save_pkl))121 122    print("ICD codebook training complete")123    print(f"dataset: {args.csv}")124    print(f"rows_total: {len(entries)}")125    print(f"saved_model_pkl: {args.save_pkl}")126 127    if args.sample_query:128        results = model.search(args.sample_query, top_k=5)129        print(f"sample_query: {args.sample_query}")130        for i, (code, score) in enumerate(results, start=1):131            print(f"top_{i}: code={code} score={score:.2f} desc={model.code_to_desc[code]}")132 133 134if __name__ == "__main__":135    main()136