structuremappingmemory 1.0.0__py3-none-any.whl

sma/eval/bugsinpy.py

@@ -0,0 +1,358 @@
+"""BugsInPy metadata parsing, patch-structure extraction, and case assembly.
+
+Blueprint T3 (bug-fix memory): a case is the deterministic structure of a
+buggy program state -- the unified diff of the fix (which files/functions
+were touched, what calls/keywords/exceptions were added or removed, how big
+the change is) plus the failing-test context from ``run_test.sh``. Everything
+here is Tier-0: regex/diff parsing only, no models.
+
+Layout of the dataset (github.com/soarsmu/BugsInPy, cloned into
+``data/raw/bugsinpy``)::
+
+    projects/<project>/bugs/<id>/bug.info        # commit ids, test_file
+    projects/<project>/bugs/<id>/bug_patch.txt   # unified diff buggy->fixed
+    projects/<project>/bugs/<id>/run_test.sh     # failing test invocation
+"""
+
+from __future__ import annotations
+
+import pathlib
+import re
+from collections import Counter
+from dataclasses import dataclass, field
+
+from sma.encoders import get_encoder
+from sma.ir.schema import Case, Statement, entity, make_case, stmt
+
+# ---------------------------------------------------------------------------
+# Discovery / loading
+# ---------------------------------------------------------------------------
+
+
+def discover_bug_metadata(root: str | pathlib.Path) -> list[pathlib.Path]:
+    return sorted(pathlib.Path(root).glob("projects/*/bugs/*/bug.info"))
+
+
+@dataclass
+class BugRecord:
+    project: str
+    bug_id: str
+    patch_text: str
+    test_files: tuple[str, ...] = ()
+    test_names: tuple[str, ...] = ()
+
+    @property
+    def key(self) -> str:
+        return f"{self.project}/{self.bug_id}"
+
+
+def _parse_bug_info(path: pathlib.Path) -> dict[str, str]:
+    info: dict[str, str] = {}
+    for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
+        m = re.match(r'^(\w+)="(.*)"\s*$', line.strip())
+        if m:
+            info[m.group(1)] = m.group(2)
+    return info
+
+
+def _parse_run_test(path: pathlib.Path) -> tuple[str, ...]:
+    """Extract failing-test function names from run_test.sh (deterministic)."""
+    if not path.exists():
+        return ()
+    names: list[str] = []
+    for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
+        line = line.strip()
+        if not line or line.startswith("#"):
+            continue
+        # pytest path/to/file.py::Class::test_name  (possibly several)
+        for token in line.split():
+            if "::" in token:
+                names.append(token.split("::")[-1])
+        # python -m unittest [-q] pkg.mod.Class.test_name
+        if "unittest" in line and "::" not in line:
+            tail = line.split()[-1]
+            comp = tail.split(".")[-1]
+            if comp.startswith("test"):
+                names.append(comp)
+    out, seen = [], set()
+    for n in names:
+        n = re.sub(r"\[.*\]$", "", n)  # strip parametrize ids
+        if n and n not in seen:
+            seen.add(n)
+            out.append(n)
+    return tuple(out)
+
+
+def load_bugs(root: str | pathlib.Path) -> list[BugRecord]:
+    """Load every bug with a non-empty patch, sorted by (project, bug id)."""
+    records: list[BugRecord] = []
+    for info_path in discover_bug_metadata(root):
+        bug_dir = info_path.parent
+        project = bug_dir.parent.parent.name
+        bug_id = bug_dir.name
+        patch_path = bug_dir / "bug_patch.txt"
+        if not patch_path.exists():
+            continue
+        patch_text = patch_path.read_text(encoding="utf-8", errors="replace")
+        if not patch_text.strip():
+            continue  # e.g. keras bug 12 ships an empty patch
+        info = _parse_bug_info(info_path)
+        test_files = tuple(
+            t.strip() for t in info.get("test_file", "").split(";") if t.strip()
+        )
+        records.append(
+            BugRecord(
+                project=project,
+                bug_id=bug_id,
+                patch_text=patch_text,
+                test_files=test_files,
+                test_names=_parse_run_test(bug_dir / "run_test.sh"),
+            )
+        )
+    records.sort(key=lambda r: (r.project, int(r.bug_id) if r.bug_id.isdigit() else 0))
+    return records
+
+
+# ---------------------------------------------------------------------------
+# Unified-diff parsing
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class Hunk:
+    file: str
+    header: str  # text after the second @@ (enclosing-scope context)
+    removed: list[str] = field(default_factory=list)
+    added: list[str] = field(default_factory=list)
+    context: list[str] = field(default_factory=list)
+
+
+@dataclass
+class PatchFacts:
+    files: list[str]
+    hunks: list[Hunk]
+    functions: list[tuple[str, str]]  # (file basename, function) pairs
+    added_lines: list[str]  # non-blank code lines added (markers stripped)
+    removed_lines: list[str]
+    exceptions: list[str]  # CamelCase *Error/*Exception/*Warning mentioned
+
+    @property
+    def n_added(self) -> int:
+        return len(self.added_lines)
+
+    @property
+    def n_removed(self) -> int:
+        return len(self.removed_lines)
+
+
+_EXC_RE = re.compile(r"\b([A-Z][A-Za-z0-9]*(?:Error|Exception|Warning))\b")
+_DEF_RE = re.compile(r"\bdef\s+([A-Za-z_]\w*)")
+
+
+def parse_patch(patch_text: str) -> PatchFacts:
+    files: list[str] = []
+    hunks: list[Hunk] = []
+    current_file = ""
+    hunk: Hunk | None = None
+    for raw in patch_text.splitlines():
+        if raw.startswith("diff --git"):
+            m = re.search(r" b/(\S+)$", raw)
+            current_file = m.group(1) if m else raw.split()[-1]
+            files.append(current_file)
+            hunk = None
+            continue
+        if raw.startswith("+++") or raw.startswith("---") or raw.startswith("index "):
+            continue
+        if raw.startswith("@@"):
+            m = re.match(r"^@@[^@]*@@\s?(.*)$", raw)
+            hunk = Hunk(file=current_file, header=m.group(1) if m else "")
+            hunks.append(hunk)
+            continue
+        if hunk is None:
+            continue
+        if raw.startswith("+"):
+            line = raw[1:]
+            if line.strip():
+                hunk.added.append(line)
+        elif raw.startswith("-"):
+            line = raw[1:]
+            if line.strip():
+                hunk.removed.append(line)
+        else:
+            line = raw[1:] if raw.startswith(" ") else raw
+            if line.strip():
+                hunk.context.append(line)
+
+    added = [l for h in hunks for l in h.added]
+    removed = [l for h in hunks for l in h.removed]
+    context = [l for h in hunks for l in h.context]
+
+    functions: list[tuple[str, str]] = []
+    seen_fn: set[tuple[str, str]] = set()
+    for h in hunks:
+        base = pathlib.PurePosixPath(h.file).name
+        names = _DEF_RE.findall(h.header)
+        # defs appearing inside the changed lines are also "modified functions"
+        for l in h.removed + h.added:
+            names.extend(_DEF_RE.findall(l))
+        if not names:
+            names = ["<module>"]
+        for name in names:
+            key = (base, name)
+            if key not in seen_fn:
+                seen_fn.add(key)
+                functions.append(key)
+
+    exceptions = sorted(
+        {e for l in added + context for e in _EXC_RE.findall(l)}
+    )
+    return PatchFacts(
+        files=files,
+        hunks=hunks,
+        functions=functions,
+        added_lines=added,
+        removed_lines=removed,
+        exceptions=exceptions,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Case assembly (Tier-0)
+# ---------------------------------------------------------------------------
+
+_PY_KEYWORDS = {
+    "if", "elif", "else", "for", "while", "try", "except", "finally", "raise",
+    "return", "with", "assert", "not", "and", "or", "in", "is", "def", "class",
+    "lambda", "yield", "del", "pass", "import", "from", "as", "global", "print",
+}
+_CALL_RE = re.compile(r"([A-Za-z_][\w\.]*)\s*\(")
+_KEYWORD_TOKENS = (
+    "if", "elif", "else", "for", "while", "try", "except", "finally", "raise",
+    "return", "with", "assert", "not", "and", "or", "in", "is", "None",
+)
+_KEYWORD_RES = {k: re.compile(rf"\b{k}\b") for k in _KEYWORD_TOKENS}
+
+_MAX_CALLS = 12
+_MAX_EXCEPTIONS = 8
+_MAX_FUNCTIONS = 16
+_MAX_ENCODER_STMTS = 24
+
+
+def _calls(lines: list[str]) -> Counter[str]:
+    counts: Counter[str] = Counter()
+    for line in lines:
+        for name in _CALL_RE.findall(line):
+            tail = name.split(".")[-1]
+            if tail in _PY_KEYWORDS or name in _PY_KEYWORDS:
+                continue
+            counts[tail] += 1
+    return counts
+
+
+def _keyword_counts(lines: list[str]) -> Counter[str]:
+    counts: Counter[str] = Counter()
+    for line in lines:
+        code = line.split("#", 1)[0]
+        for kw, rx in _KEYWORD_RES.items():
+            counts[kw] += len(rx.findall(code))
+    return counts
+
+
+def size_bucket(n: int) -> str:
+    if n <= 0:
+        return "0"
+    if n == 1:
+        return "1"
+    if n <= 3:
+        return "2-3"
+    if n <= 7:
+        return "4-7"
+    if n <= 15:
+        return "8-15"
+    if n <= 31:
+        return "16-31"
+    return "32plus"
+
+
+def _encoder_statements(added_lines: list[str]) -> list[Statement]:
+    """Run the existing code adapter on the added hunks (AST if it parses,
+    regex fallback otherwise); keep its non-placeholder statements."""
+    block = "\n".join(l.strip() for l in added_lines)
+    if not block.strip():
+        return []
+    encoder = get_encoder("code")
+    result = encoder.encode(block, language="python")
+    functors = {s.functor for s in result.case.statements}
+    if "syntaxError" in functors:  # hunk fragments rarely parse; regex fallback
+        result = encoder.encode(block, language="text")
+    keep = [
+        s
+        for s in result.case.statements
+        if s.functor not in {"syntaxError", "emptyCode", "rawCode"}
+    ]
+    return keep[:_MAX_ENCODER_STMTS]
+
+
+def bug_case(record: BugRecord, facts: PatchFacts | None = None) -> Case:
+    """Deterministic Tier-0 case for one bug: diff structure + failing test."""
+    facts = facts or parse_patch(record.patch_text)
+    statements: set[Statement] = set()
+
+    for base, fn in facts.functions[:_MAX_FUNCTIONS]:
+        statements.add(stmt("modifies", entity(base, "file"), entity(fn, "function")))
+    statements.add(stmt("addsLines", entity(size_bucket(facts.n_added), "count")))
+    statements.add(stmt("removesLines", entity(size_bucket(facts.n_removed), "count")))
+    statements.add(
+        stmt("touchesFiles", entity(size_bucket(len(facts.files)), "count"))
+    )
+    if any("test" in f.lower() for f in facts.files):
+        statements.add(stmt("touchesTestFile", entity("patch", "scope")))
+
+    for exc in facts.exceptions[:_MAX_EXCEPTIONS]:
+        statements.add(stmt("mentionsException", entity(exc, "exception")))
+
+    added_calls = _calls(facts.added_lines)
+    removed_calls = _calls(facts.removed_lines)
+    net_added = sorted(
+        (c for c in added_calls if added_calls[c] > removed_calls.get(c, 0)),
+        key=lambda c: (-added_calls[c], c),
+    )[:_MAX_CALLS]
+    net_removed = sorted(
+        (c for c in removed_calls if removed_calls[c] > added_calls.get(c, 0)),
+        key=lambda c: (-removed_calls[c], c),
+    )[:_MAX_CALLS]
+    for name in net_added:
+        statements.add(stmt("addsCall", entity("patch", "scope"), entity(name, "callable")))
+    for name in net_removed:
+        statements.add(stmt("removesCall", entity("patch", "scope"), entity(name, "callable")))
+
+    kw_added = _keyword_counts(facts.added_lines)
+    kw_removed = _keyword_counts(facts.removed_lines)
+    for kw in _KEYWORD_TOKENS:
+        if kw_added[kw] > kw_removed[kw]:
+            statements.add(stmt("addsKeyword", entity("patch", "scope"), entity(kw, "keyword")))
+        elif kw_removed[kw] > kw_added[kw]:
+            statements.add(stmt("removesKeyword", entity("patch", "scope"), entity(kw, "keyword")))
+
+    for name in record.test_names[:4]:
+        statements.add(stmt("failingTest", entity(name, "test")))
+    for tf in record.test_files[:4]:
+        statements.add(
+            stmt("testModule", entity(pathlib.PurePosixPath(tf).name, "file"))
+        )
+
+    statements.update(_encoder_statements(facts.added_lines))
+
+    if not statements:
+        statements.add(stmt("emptyPatch", entity(record.key, "bug")))
+    return make_case(
+        sorted(statements, key=repr),
+        {
+            "adapter": "code",
+            "tier": 0,
+            "dataset": "bugsinpy",
+            "project": record.project,
+            "bug": record.bug_id,
+        },
+        case_id=f"bugsinpy:{record.key}",
+    )

sma/eval/bugsinpy_families.py

@@ -0,0 +1,164 @@
+"""Deterministic fix-pattern categories for BugsInPy patches.
+
+Categories are assigned by ORDERED rules over the unified diff (first match
+wins). They are the ground-truth labels for the T3 retrieval metric
+(fix-category-hit@k): did retrieval surface a past bug fixed by the SAME
+KIND of change, not just a textually similar file?
+
+Rule order (per the T3 specification):
+  1. add-null-check          adds ``is None`` / ``is not None`` / ``if not x``
+  2. exception-handling      adds ``try:`` / ``except`` / ``raise``
+  3. boundary                comparison-operator swap or +/-1 on an otherwise
+                             identical line (off-by-one / boundary fixes)
+  4. type-coercion           adds a builtin cast (int()/str()/list()/...)
+  5. api-substitution        small hunk replaces one call with another
+  6. default-arg-change      a ``def`` signature line changes its defaults/args
+  7. condition-strengthening an if/elif/while gains and/or clauses
+  8. other
+
+"Net added" means the pattern occurs more often in added lines than in
+removed lines, so pure code motion does not trigger a rule.
+"""
+
+from __future__ import annotations
+
+import re
+
+from .bugsinpy import Hunk, PatchFacts, _calls
+
+CATEGORIES = (
+    "add-null-check",
+    "exception-handling",
+    "boundary",
+    "type-coercion",
+    "api-substitution",
+    "default-arg-change",
+    "condition-strengthening",
+    "other",
+)
+
+_NULL_CHECK = re.compile(r"\bis\s+(?:not\s+)?None\b|^\s*(?:el)?if\s+not\s+\w")
+_EXC_HANDLING = re.compile(r"^\s*(?:try\s*:|except[\s:(]|raise\b)")
+_CAST = re.compile(
+    r"(?<![\w.])(?:int|str|float|bool|list|tuple|set|dict|frozenset|bytes)\(|\.astype\("
+)
+_CMP_OPS = re.compile(r"<=|>=|==|!=|<|>")
+_PM_ONE = re.compile(r"[+-]\s*1\b")
+_COND_LINE = re.compile(r"^\s*((?:el)?if|while)\b")
+_BOOL_OP = re.compile(r"\b(?:and|or)\b")
+_DEF_LINE = re.compile(r"^\s*def\s+([A-Za-z_]\w*)\s*\((.*)$")
+
+
+def _strip_code(line: str) -> str:
+    return line.split("#", 1)[0]
+
+
+def _net(pattern: re.Pattern[str], added: list[str], removed: list[str]) -> bool:
+    n_add = sum(len(pattern.findall(_strip_code(l))) for l in added)
+    n_rem = sum(len(pattern.findall(_strip_code(l))) for l in removed)
+    return n_add > n_rem
+
+
+def _paired_swap(hunk: Hunk, pattern: re.Pattern[str]) -> bool:
+    """True if some removed/added line pair is identical once `pattern`
+    occurrences are masked out, while the raw lines differ (i.e. the ONLY
+    change is in the matched operator/constant)."""
+    def norm(line: str) -> str:
+        return re.sub(r"\s+", " ", pattern.sub("\x00", _strip_code(line))).strip()
+
+    removed = {norm(l): _strip_code(l).strip() for l in hunk.removed if pattern.search(_strip_code(l))}
+    for line in hunk.added:
+        code = _strip_code(line)
+        if not pattern.search(code):
+            continue
+        key = norm(line)
+        if key in removed and removed[key] != code.strip() and "\x00" in key:
+            return True
+    return False
+
+
+def _is_null_check(facts: PatchFacts) -> bool:
+    return _net(_NULL_CHECK, facts.added_lines, facts.removed_lines)
+
+
+def _is_exception_handling(facts: PatchFacts) -> bool:
+    return _net(_EXC_HANDLING, facts.added_lines, facts.removed_lines)
+
+
+def _is_boundary(facts: PatchFacts) -> bool:
+    for hunk in facts.hunks:
+        if _paired_swap(hunk, _CMP_OPS) or _paired_swap(hunk, _PM_ONE):
+            return True
+    return False
+
+
+def _is_type_coercion(facts: PatchFacts) -> bool:
+    return _net(_CAST, facts.added_lines, facts.removed_lines)
+
+
+def _is_api_substitution(facts: PatchFacts) -> bool:
+    for hunk in facts.hunks:
+        if not (1 <= len(hunk.removed) <= 3 and 1 <= len(hunk.added) <= 3):
+            continue
+        removed_calls = set(_calls(hunk.removed))
+        added_calls = set(_calls(hunk.added))
+        if (removed_calls - added_calls) and (added_calls - removed_calls):
+            return True
+    return False
+
+
+def _is_default_arg_change(facts: PatchFacts) -> bool:
+    for hunk in facts.hunks:
+        removed_defs = {}
+        for line in hunk.removed:
+            m = _DEF_LINE.match(_strip_code(line))
+            if m:
+                removed_defs[m.group(1)] = m.group(2).strip()
+        for line in hunk.added:
+            m = _DEF_LINE.match(_strip_code(line))
+            if not m:
+                continue
+            name, args = m.group(1), m.group(2).strip()
+            old = removed_defs.get(name)
+            if old is not None and old != args and ("=" in old or "=" in args):
+                return True
+    return False
+
+
+def _is_condition_strengthening(facts: PatchFacts) -> bool:
+    for hunk in facts.hunks:
+        removed_conds: dict[str, int] = {}
+        for line in hunk.removed:
+            m = _COND_LINE.match(_strip_code(line))
+            if m:
+                kw = m.group(1)
+                n = len(_BOOL_OP.findall(_strip_code(line)))
+                removed_conds[kw] = max(removed_conds.get(kw, -1), n)
+        for line in hunk.added:
+            m = _COND_LINE.match(_strip_code(line))
+            if not m:
+                continue
+            kw = m.group(1)
+            n = len(_BOOL_OP.findall(_strip_code(line)))
+            if kw in removed_conds and n > removed_conds[kw]:
+                return True
+    return False
+
+
+_RULES = (
+    ("add-null-check", _is_null_check),
+    ("exception-handling", _is_exception_handling),
+    ("boundary", _is_boundary),
+    ("type-coercion", _is_type_coercion),
+    ("api-substitution", _is_api_substitution),
+    ("default-arg-change", _is_default_arg_change),
+    ("condition-strengthening", _is_condition_strengthening),
+)
+
+
+def categorize(facts: PatchFacts) -> str:
+    """Ordered first-match-wins fix-pattern category for one patch."""
+    for name, rule in _RULES:
+        if rule(facts):
+            return name
+    return "other"

sma/eval/crossdomain.py

@@ -0,0 +1,89 @@
+"""Shared cross-domain arm evaluator (4b).
+
+Guarantees an IDENTICAL index/query partition across arms (generic / drafted /
+expert) by splitting on the original record order — NOT on the encoding-dependent
+case_id — so the baselines are stable and the only thing that varies between arms
+is the SMA encoding. Returns the result dict and writes cd_<domain>_<phase>.csv.
+"""
+from __future__ import annotations
+
+import csv
+import pathlib
+import time
+
+from sma.index.macfac import MacFacIndex
+from sma.ir.schema import Statement
+from sma.eval.baselines.bm25 import rank_bm25_like
+from sma.eval.baselines.dense import rank_tfidf_dense_batch
+from sma.eval.metrics import macro_f1
+from sma.eval.stats import paired_bootstrap, holm_bonferroni, cliffs_delta
+
+OUT = pathlib.Path("reports/confirmatory")
+
+
+def _ho_density(case) -> float:
+    ex = case.expressions()
+    ho = sum(1 for s in ex if any(isinstance(a, Statement) for a in s.args))
+    return ho / max(len(ex), 1)
+
+
+def _vote(ranked_ids, label_of, labels):
+    pos, neg = labels
+    tally = {pos: 0, neg: 0}
+    for cid in ranked_ids:
+        tally[label_of[cid]] += 1
+    return pos if tally[pos] >= tally[neg] else neg
+
+
+def evaluate_arm(items, encode_fn, row_text_fn, labels, phase, domain,
+                 k=10, frac=0.7):
+    """items: list of records (each has .label). encode_fn(item)->Case.
+    row_text_fn(item)->str for the baselines. labels=(pos,neg)."""
+    index_n = int(len(items) * frac)
+    index_items, query_items = items[:index_n], items[index_n:]   # FIXED split
+
+    def enc(its):
+        return [(it, encode_fn(it)) for it in its]
+    idx, qry = enc(index_items), enc(query_items)
+    label_of = {c.case_id: it.label for it, c in idx + qry}
+    text_of = {c.case_id: row_text_fn(it) for it, c in idx + qry}
+    dens = [_ho_density(c) for _, c in idx + qry]
+    mean_ho = sum(dens) / len(dens)
+
+    index_cases = [c for _, c in idx]
+    index_docs = [(c.case_id, text_of[c.case_id]) for _, c in idx]
+    t0 = time.perf_counter()
+    sma_index = MacFacIndex(); sma_index.build(index_cases)
+    gold, sma_p, bm_p, dn_p = [], [], [], []
+    dense_rk = rank_tfidf_dense_batch([text_of[c.case_id] for _, c in qry], index_docs, k=k)
+    for qi, (_, qc) in enumerate(qry):
+        gold.append(label_of[qc.case_id])
+        res = sma_index.retrieve(qc, k=k, shortlist=60, fac_budget=25)
+        sma_p.append(_vote([r.case_id for r in res], label_of, labels))
+        bm = rank_bm25_like(text_of[qc.case_id], index_docs, k=k)
+        bm_p.append(_vote([cid for cid, _ in bm], label_of, labels))
+        dn_p.append(_vote([cid for cid, _ in dense_rk[qi]], label_of, labels))
+    dt = time.perf_counter() - t0
+
+    f1 = {"SMA": macro_f1(gold, sma_p), "BM25": macro_f1(gold, bm_p), "Dense": macro_f1(gold, dn_p)}
+    print(f"[{domain}/{phase}] HO-density={mean_ho:.4f}  macro-F1: SMA {f1['SMA']:.4f}  "
+          f"BM25 {f1['BM25']:.4f}  Dense {f1['Dense']:.4f}  ({dt:.0f}s)", flush=True)
+
+    def correct(p):
+        return [1.0 if a == g else 0.0 for a, g in zip(p, gold)]
+    sma_c = correct(sma_p); pv, summ = {}, []
+    for name, pred in (("BM25", bm_p), ("Dense", dn_p)):
+        bs = paired_bootstrap(sma_c, correct(pred)); pv[name] = bs["p_value"]
+        summ.append({"phase": phase, "domain": domain, "baseline": name,
+                     "ho_density": f"{mean_ho:.4f}", "sma_f1": f"{f1['SMA']:.4f}",
+                     "baseline_f1": f"{f1[name]:.4f}", "delta": f"{bs['delta']:.4f}",
+                     "ci_low": f"{bs['ci_low']:.4f}", "ci_high": f"{bs['ci_high']:.4f}",
+                     "cliffs": f"{cliffs_delta(sma_c, correct(pred)):.4f}"})
+    holm = holm_bonferroni(pv)
+    for s in summ:
+        s["p_holm"] = f"{holm[s['baseline']]:.4f}"
+    out = OUT / f"cd_{domain}_{phase}.csv"
+    out.parent.mkdir(parents=True, exist_ok=True)
+    with out.open("w", newline="") as fh:
+        w = csv.DictWriter(fh, fieldnames=list(summ[0])); w.writeheader(); w.writerows(summ)
+    return {"f1": f1, "ho": mean_ho, "summ": summ}

sma/eval/diabetes.py

@@ -0,0 +1,61 @@
+"""Diabetes-130 loader + per-encounter artifact builder (4b cross-domain).
+
+Real UCI EHR. We build a flat CSV-row artifact for the GENERIC structured
+adapter (the honest 'before': flat triples, no higher-order relations) and a
+plain attr=value text for the lexical/dense baselines. The readmission label is
+NEVER encoded (leakage guard).
+"""
+from __future__ import annotations
+
+import csv
+import pathlib
+import random
+from dataclasses import dataclass
+
+# ids, mostly-missing columns, and the LABEL are excluded from the encoding.
+DROP = {"encounter_id", "patient_nbr", "weight", "payer_code", "readmitted"}
+
+
+@dataclass(frozen=True)
+class Encounter:
+    eid: str
+    fields: dict[str, str]   # cleaned attribute -> value (no '?'/'' )
+    label: str               # "early" (readmitted <30 days) vs "not"
+
+
+def _csv_path() -> pathlib.Path:
+    return next(pathlib.Path("data/raw/diabetes130").rglob("diabetic_data.csv"))
+
+
+def load_encounters(sample: int = 1500, seed: int = 7, balanced: bool = True) -> list[Encounter]:
+    rows = list(csv.DictReader(_csv_path().open()))
+    rng = random.Random(seed)
+    rng.shuffle(rows)
+
+    def to_enc(r: dict) -> Encounter:
+        fields = {k: v for k, v in r.items() if k not in DROP and v not in ("?", "")}
+        label = "early" if r["readmitted"] == "<30" else "not"
+        return Encounter(r["encounter_id"], fields, label)
+
+    encs = [to_enc(r) for r in rows]
+    if not balanced:
+        return encs[:sample]
+    # balance the two classes so retrieval-by-analogy has signal (early is ~11%)
+    early = [e for e in encs if e.label == "early"]
+    notr = [e for e in encs if e.label == "not"]
+    half = sample // 2
+    out = early[:half] + notr[:half]
+    rng.shuffle(out)
+    return out
+
+
+def row_csv(enc: Encounter) -> str:
+    """Two-line CSV (header + one row) for the structured adapter -> flat
+    (attribute row value) triples."""
+    keys = sorted(enc.fields)
+    return ",".join(keys) + "\n" + ",".join(enc.fields[k] for k in keys) + "\n"
+
+
+def row_text(enc: Encounter) -> str:
+    """attr=value text for BM25 / dense baselines."""
+    return " ".join(f"{k}={v}" for k, v in sorted(enc.fields.items()))