natural-memory-nm21/train_attribute_head.py

"""Train the attribute-coverage head: "which attribute is this question asking about?"

Why this exists
---------------
The runtime answers ~75% of unanswerable paraphrased questions with an invented value.
A score threshold cannot fix that: measured on the same frozen keys, the best fitted head
over nine score-geometry features reaches AUC 0.61 (and ~0.5 for three of four scorers),
because 24 same-shape candidates look equally plausible for any question.

What works instead is a different question: *which attribute is being asked about, and does
the bank actually hold it?*  This script trains that head and saves it as a deployable
artifact.

Labels are recovered exactly, not guessed: every query in this dataset is one of the
authored paraphrases, so the asked-about attribute is looked up from the paraphrase table
rather than taken from the episode metadata (which records the answerable *target* and is
therefore wrong for abstention episodes).

Usage::

    python -m V2_dpskw.train_attribute_head --data-dir data/zero_overlap ^
        --feature-cache H:\\Memory\\nm_cache\\nm_zero_overlap\\feature_cache ^
        --output-dir checkpoints/memory_attribute_head
"""

from __future__ import annotations

import argparse
import hashlib
import json
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
import torch

from .make_zero_overlap_paraphrase_data import ATTRIBUTE_PARAPHRASES


def text_key(text: str) -> str:
    return hashlib.sha1(str(text).encode("utf-8", "replace")).hexdigest()


def paraphrase_table() -> dict[str, str]:
    table = {}
    for attribute, paraphrases in ATTRIBUTE_PARAPHRASES:
        for paraphrase in paraphrases:
            table[paraphrase] = attribute
    return table


def load(path: Path) -> list[dict]:
    rows = []
    with path.open("r", encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def softmax_fit(X: np.ndarray, y: np.ndarray, classes: int, *, steps: int = 6000,
                lr: float = 0.5, l2: float = 1e-3):
    W = np.zeros((X.shape[1], classes))
    b = np.zeros(classes)
    n = len(y)
    onehot = np.zeros((n, classes))
    onehot[np.arange(n), y] = 1.0
    for _ in range(steps):
        logits = X @ W + b
        logits -= logits.max(axis=1, keepdims=True)
        probs = np.exp(logits)
        probs /= probs.sum(axis=1, keepdims=True)
        grad = probs - onehot
        W -= lr * (X.T @ grad / n + l2 * W)
        b -= lr * grad.mean(axis=0)
    return W, b


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--data-dir", default="data/zero_overlap")
    parser.add_argument("--feature-cache", default=r"H:\Memory\nm_cache\nm_zero_overlap\feature_cache")
    parser.add_argument("--output-dir", default="checkpoints/memory_attribute_head")
    parser.add_argument("--report", default="attribute_head_training.json")
    args = parser.parse_args()

    cache = Path(args.feature_cache)
    bank = np.load(cache / "features.f16.npy", mmap_mode="r")
    lookup = json.loads((cache / "index.json").read_text(encoding="utf-8"))
    table = paraphrase_table()
    attributes = sorted({name for name, _ in ATTRIBUTE_PARAPHRASES})
    index_of = {name: position for position, name in enumerate(attributes)}

    def build(rows):
        X, y, rows_used = [], [], []
        for row in rows:
            attribute = table.get(row["query"])
            if attribute is None:
                raise SystemExit(f"query is not one of the authored paraphrases: {row['query']!r}")
            key = text_key(row["query"])
            if key not in lookup:
                raise SystemExit(f"query missing from the feature bank: {row['query']!r}")
            X.append(np.asarray(bank[lookup[key]], dtype=np.float32))
            y.append(index_of[attribute])
            rows_used.append(row)
        return np.array(X), np.array(y), rows_used

    train_X, train_y, _ = build(load(Path(args.data_dir) / "train.jsonl"))
    eval_X, eval_y, eval_rows = build(load(Path(args.data_dir) / "eval.jsonl"))

    train_Xn = train_X / (np.linalg.norm(train_X, axis=1, keepdims=True) + 1e-6)
    eval_Xn = eval_X / (np.linalg.norm(eval_X, axis=1, keepdims=True) + 1e-6)
    W, b = softmax_fit(train_Xn, train_y, len(attributes))

    def probabilities(X):
        logits = X @ W + b
        logits -= logits.max(axis=1, keepdims=True)
        p = np.exp(logits)
        return p / p.sum(axis=1, keepdims=True)

    eval_probs = probabilities(eval_Xn)
    eval_pred = eval_probs.argmax(axis=1)
    answerable = np.array([bool(row.get("positive_indices")) for row in eval_rows])

    # Coverage = the asked-about attribute is actually present in this episode's bank.
    predicted_names = [attributes[position] for position in eval_pred]
    present_sets = [{c.get("attribute") for c in row["candidates"] if c.get("attribute")}
                    for row in eval_rows]
    covered = np.array([name in present for name, present in zip(predicted_names, present_sets)])

    report = {
        "built_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "attributes": attributes,
        "classes": len(attributes),
        "chance_pct": round(100.0 / len(attributes), 2),
        "train_queries": int(train_y.size),
        "eval_queries": int(eval_y.size),
        "attribute_accuracy_train_pct": round(100 * float(np.mean(
            probabilities(train_Xn).argmax(axis=1) == train_y)), 2),
        "attribute_accuracy_eval_pct": round(100 * float(np.mean(eval_pred == eval_y)), 2),
        "attribute_accuracy_eval_answerable_pct": round(100 * float(np.mean(
            (eval_pred == eval_y)[answerable])), 2),
        "coverage_rule": {
            "answerable_covered_pct": round(100 * float(np.mean(covered[answerable])), 2),
            "abstention_covered_pct": round(100 * float(np.mean(covered[~answerable])), 2),
            "unknown_refusal_pct": round(100 * float(np.mean(~covered[~answerable])), 2),
            "known_false_refusal_pct": round(100 * float(np.mean(~covered[answerable])), 2),
            "counts": {
                "tp": int(np.sum(~covered & ~answerable)),
                "fn": int(np.sum(covered & ~answerable)),
                "fp": int(np.sum(~covered & answerable)),
                "tn": int(np.sum(covered & answerable)),
            },
        },
    }

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    torch.save({
        "weight": torch.tensor(W, dtype=torch.float32),
        "bias": torch.tensor(b, dtype=torch.float32),
        "attributes": attributes,
        "format_version": 1,
        "input": "l2_normalised frozen query key (hidden_size)",
    }, out_dir / "attribute_head.pt")
    (out_dir / "attribute_head_meta.json").write_text(
        json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    Path(args.report).write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(report, ensure_ascii=False, indent=2), flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())