qarar 1.0.2__py3-none-any.whl

nomeda/__init__.py

@@ -0,0 +1,71 @@
+"""NOMEDA-Decision v1: one Arabic state -> many independent typed decisions."""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from .config import (NomedaConfig, NomedConfig, DEFAULT_CONFIG, k_bucket,
+                     K_BUCKETS, K_BUCKET_NAMES)
+from .model import NomedaModel, NomedModel
+from .runtime import NomedaRuntime, NomedRuntime, validate_questions
+
+__version__ = "1.0.0"
+__all__ = ["load", "Nomeda", "Nomed",
+           "NomedaConfig", "NomedaModel", "NomedaRuntime",
+           "NomedConfig", "NomedModel", "NomedRuntime",
+           "validate_questions", "k_bucket"]
+
+
+class Nomeda:
+    """Loaded model + calibration; ``predict`` mirrors the Laya output shape."""
+
+    def __init__(self, model: NomedaModel, temperatures: dict | None = None,
+                 thresholds: dict | None = None, device: str = "cpu"):
+        self.model = model
+        self.encoder_id = model.cfg.encoder_id
+        self.encoder_revision = model.cfg.encoder_revision
+        self.runtime = NomedaRuntime(model, temperatures or {}, thresholds or {}, device=device)
+
+    def predict(self, state, questions):
+        return self.runtime.predict(state, questions)
+
+    def predict_batch(self, states, questions):
+        return self.runtime.predict_batch(states, questions)
+
+    def param_report(self) -> dict:
+        return self.model.param_report()
+
+
+def load(path_or_hub_id: str, device: str = "auto", tiny: bool = False) -> Nomeda:
+    """Load a NOMEDA-Decision model from a local export directory or a Hub id."""
+    if device == "auto":
+        import torch
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+    if Path(path_or_hub_id).is_dir():
+        root = Path(path_or_hub_id)
+        raw = json.loads((root / "config.json").read_text(encoding="utf-8"))
+        cfg = NomedaConfig(**{k: v for k, v in raw.items() if k in NomedaConfig.__dataclass_fields__})
+        enc_dir = root / "encoder"
+        model = NomedaModel(cfg, device=device, tiny=tiny)
+        if enc_dir.exists():
+            from transformers import AutoModel, AutoTokenizer
+            model.encoder.tok = AutoTokenizer.from_pretrained(enc_dir)
+            model.encoder.model = AutoModel.from_pretrained(enc_dir).to(device)
+        head = root / "decision_head.safetensors"
+        if head.exists():
+            from safetensors.torch import load_file
+            state = load_file(str(head))
+            model.load_state_dict(state, strict=False)
+        temps = json.loads((root / "temperatures.json").read_text(encoding="utf-8")) if (root / "temperatures.json").exists() else {}
+        thr = json.loads((root / "thresholds.json").read_text(encoding="utf-8")) if (root / "thresholds.json").exists() else {}
+        model.eval()
+        model.to(device)
+        return Nomeda(model, temps, thr, device=device)
+    # Hub id
+    from huggingface_hub import snapshot_download
+    d = snapshot_download(repo_id=path_or_hub_id)
+    return load(d, device=device, tiny=tiny)
+
+
+# ---- backward-compatible alias (published as `Nomed` in v0.1-beta) --------
+Nomed = Nomeda

nomeda/blocks.py

@@ -0,0 +1,112 @@
+"""Decision blocks, option-state block and attention pooling (§3.4–3.6)."""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MultiheadAttention(nn.Module):
+    def __init__(self, d: int, n_heads: int, dropout: float = 0.0):
+        super().__init__()
+        assert d % n_heads == 0
+        self.d, self.h = d, n_heads
+        self.dh = d // n_heads
+        self.q = nn.Linear(d, d)
+        self.k = nn.Linear(d, d)
+        self.v = nn.Linear(d, d)
+        self.o = nn.Linear(d, d)
+        self.dropout = dropout
+
+    def forward(self, q, k, v, q_mask=None, k_mask=None):
+        B, Lq, _ = q.shape
+        Lk = k.shape[1]
+        Q = self.q(q).view(B, Lq, self.h, self.dh).transpose(1, 2)
+        K = self.k(k).view(B, Lk, self.h, self.dh).transpose(1, 2)
+        V = self.v(v).view(B, Lk, self.h, self.dh).transpose(1, 2)
+        attn_mask = None
+        if k_mask is not None:
+            attn_mask = torch.zeros(B, 1, 1, Lk, device=q.device, dtype=Q.dtype)
+            attn_mask = attn_mask.masked_fill(~k_mask[:, None, None, :].bool(), -1e4)
+        out = F.scaled_dot_product_attention(
+            Q, K, V, attn_mask=attn_mask,
+            dropout_p=self.dropout if self.training else 0.0)
+        out = out.transpose(1, 2).reshape(B, Lq, self.d)
+        return self.o(out)
+
+
+class FFN(nn.Module):
+    def __init__(self, d: int, mult: int, dropout: float):
+        super().__init__()
+        self.net = nn.Sequential(nn.Linear(d, d * mult), nn.GELU(), nn.Dropout(dropout), nn.Linear(d * mult, d))
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class DecisionBlock(nn.Module):
+    """Self-attn within a question -> cross-attn to its own state -> FFN (pre-norm)."""
+
+    def __init__(self, d, n_heads, mult, dropout):
+        super().__init__()
+        self.ln1, self.ln2, self.ln3 = nn.LayerNorm(d), nn.LayerNorm(d), nn.LayerNorm(d)
+        self.self_attn = MultiheadAttention(d, n_heads, dropout)
+        self.cross_attn = MultiheadAttention(d, n_heads, dropout)
+        self.ffn = FFN(d, mult, dropout)
+
+    def forward(self, Q, q_mask, S_pad, S_mask, state_idx):
+        Q = Q + self.self_attn(self.ln1(Q), self.ln1(Q), self.ln1(Q), q_mask=q_mask, k_mask=q_mask)
+        # cross-attention grouped by state: gather each question's own state (vectorised).
+        # Equivalent to a per-state loop but without one kernel launch + sync per state.
+        x = self.ln2(Q)
+        k = S_pad.index_select(0, state_idx)
+        kmask = S_mask.index_select(0, state_idx)
+        Q = Q + self.cross_attn(x, k, k, k_mask=kmask)
+        Q = Q + self.ffn(self.ln3(Q))
+        return Q
+
+
+class OptionStateBlock(nn.Module):
+    def __init__(self, d, n_heads, mult, dropout):
+        super().__init__()
+        self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
+        self.cross_attn = MultiheadAttention(d, n_heads, dropout)
+        self.ffn = FFN(d, mult, dropout)
+
+    def forward(self, E, S_pad, S_mask, state_idx):
+        x = self.ln1(E).unsqueeze(1)  # [Nopt,1,d]
+        k = S_pad.index_select(0, state_idx)
+        kmask = S_mask.index_select(0, state_idx)
+        out = self.cross_attn(x, k, k, k_mask=kmask)
+        E = E + out.squeeze(1)
+        E = E + self.ffn(self.ln2(E))
+        return E
+
+
+class AttnPool(nn.Module):
+    def __init__(self, d):
+        super().__init__()
+        self.query = nn.Parameter(torch.randn(d) * 0.02)
+        self.ln = nn.LayerNorm(d)
+
+    def forward(self, Q, mask):
+        attn = (Q @ self.query) / (Q.shape[-1] ** 0.5)  # [B,L]
+        attn = attn.masked_fill(~mask.bool(), -1e4)
+        p = torch.softmax(attn, dim=-1).unsqueeze(1)
+        pooled = torch.bmm(p, Q).squeeze(1)
+        return self.ln(pooled)
+
+
+class CompatibilityScorer(nn.Module):
+    def __init__(self, d, dropout):
+        super().__init__()
+        self.ln = nn.LayerNorm(4 * d)
+        self.w1 = nn.Linear(4 * d, d)
+        self.w2 = nn.Linear(d, 1)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, q, E):
+        h = torch.cat([q, E, q * E, (q - E).abs()], dim=-1)
+        h = self.ln(h)
+        h = self.dropout(F.gelu(self.w1(h)))
+        return self.w2(h).squeeze(-1)

nomeda/calibrate.py

@@ -0,0 +1,165 @@
+"""Calibration: temperatures per (type, K-bucket) and abstention thresholds (§6)."""
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+from .config import k_bucket
+
+
+def ece(probs: List[float], correct: List[int], n_bins: int = 15) -> float:
+    """Expected calibration error, confidence = max(p), equal-width bins (Laya's definition)."""
+    if not probs:
+        return 0.0
+    total = len(probs)
+    out = 0.0
+    for b in range(n_bins):
+        lo, hi = b / n_bins, (b + 1) / n_bins
+        idx = [i for i, p in enumerate(probs) if (p > lo or (b == 0 and p >= lo)) and p <= hi]
+        if not idx:
+            continue
+        conf = sum(probs[i] for i in idx) / len(idx)
+        acc = sum(correct[i] for i in idx) / len(idx)
+        out += (len(idx) / total) * abs(conf - acc)
+    return out
+
+
+def fit_temperature(logits_list, targets_list, clamp: Tuple[float, float] = (0.5, 5.0)) -> float:
+    """Minimise NLL over a single temperature; logits_list holds per-item 1-D tensors
+    (K may vary between items)."""
+    logT = torch.zeros(1, requires_grad=True)
+    opt = torch.optim.LBFGS([logT], lr=0.1, max_iter=50)
+    tg = [int(t) for t in targets_list]
+
+    def closure():
+        opt.zero_grad()
+        T = torch.exp(logT).clamp(*clamp)
+        loss = 0.0
+        for l, t in zip(logits_list, tg):
+            loss = loss + F.cross_entropy((l / T).unsqueeze(0).float(), torch.tensor([t]))
+        loss = loss / max(1, len(logits_list))
+        loss.backward()
+        return loss
+
+    opt.step(closure)
+    with torch.no_grad():
+        return float(torch.exp(logT).clamp(*clamp))
+
+
+def fit_temperatures(buckets: Dict[Tuple[str, str], Tuple[list, list]],
+                     min_items: int = 200) -> dict:
+    """buckets: {(type, K-bucket): ([logits_1d...], [targets...])} -> temperatures.json dict."""
+    out: Dict[str, Dict[str, float]] = {}
+    per_type: Dict[str, Tuple[list, list]] = {}
+    for (qtype, kb), (logits, targets) in buckets.items():
+        per_type.setdefault(qtype, ([], []))
+        per_type[qtype][0].extend(logits)
+        per_type[qtype][1].extend(targets)
+        if len(logits) < min_items:
+            continue
+        out.setdefault(qtype, {})[kb] = fit_temperature(logits, targets)
+    for qtype, (ls, ts) in per_type.items():
+        out.setdefault(qtype, {})["default"] = fit_temperature(ls, ts)
+    return out
+
+
+def fit_abstention(conf: List[float], correct: List[int], target_acc: float = 0.95) -> float:
+    """Smallest threshold so answered questions reach accuracy >= target_acc."""
+    pairs = sorted(zip(conf, correct), reverse=True)
+    if not pairs:
+        return 0.0
+    best_t = 1.0
+    n_ok = 0
+    for i, (c, ok) in enumerate(pairs):
+        n_ok += ok
+        acc = n_ok / (i + 1)
+        if acc >= target_acc:
+            best_t = c
+    return float(best_t)
+
+
+def coverage(conf: List[float], threshold: float) -> float:
+    if not conf:
+        return 0.0
+    return sum(1 for c in conf if c >= threshold) / len(conf)
+
+
+def save(path: str, obj: dict) -> None:
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(obj, f, ensure_ascii=False, indent=2)
+
+
+# ---- fair calibrated comparison (cal_cv, §0.4) ---------------------------
+def group_of(record_id, view_group) -> str:
+    """Fold-group id: view_group and V2 en↔ar pairs stay together."""
+    if view_group:
+        return f"vg:{view_group}"
+    rid = str(record_id)
+    if ":en:" in rid or ":ar:" in rid:
+        return "pair:" + rid.replace(":en:", ":*:").replace(":ar:", ":*:")
+    return f"solo:{rid}"
+
+
+def _fold_of(group_id: str, folds: int) -> int:
+    h = int(hashlib.sha1(str(group_id).encode()).hexdigest()[:8], 16)
+    return h % folds
+
+
+def _fit_temps_from_items(items: List[dict], min_items: int,
+                          clamp: Tuple[float, float] = (0.5, 5.0)) -> dict:
+    buckets: Dict[Tuple[str, str], Tuple[list, list]] = defaultdict(lambda: ([], []))
+    per_type: Dict[str, Tuple[list, list]] = defaultdict(lambda: ([], []))
+    for it in items:
+        kb = it["kb"]
+        buckets[(it["type"], kb)][0].append(it["logits"])
+        buckets[(it["type"], kb)][1].append(it["idx"])
+        per_type[it["type"]][0].append(it["logits"])
+        per_type[it["type"]][1].append(it["idx"])
+    out: Dict[str, dict] = {}
+    for (qt, kb), (L, T) in buckets.items():
+        if len(L) >= min_items:
+            out.setdefault(qt, {})[kb] = fit_temperature(L, T, clamp)
+    for qt, (L, T) in per_type.items():
+        out.setdefault(qt, {})["default"] = fit_temperature(L, T, clamp)
+    return out
+
+
+def _score_items(items: List[dict], temps: dict) -> Tuple[float, float, List[float], List[int]]:
+    nll = 0.0
+    confs, oks = [], []
+    for it in items:
+        T = float(temps.get(it["type"], {}).get(it["kb"], temps.get(it["type"], {}).get("default", 1.0)))
+        p = torch.softmax(it["logits"] / T, dim=-1)
+        nll += -math.log(max(float(p[it["idx"]]), 1e-12))
+        confs.append(float(p.max()))
+        oks.append(int(int(torch.argmax(p)) == it["idx"]))
+    return nll / max(1, len(items)), ece(confs, oks), confs, oks
+
+
+def cal_cv_from_rows(rows: List[dict], folds: int = 2, min_items: int = 100) -> dict:
+    """Fit on fold A / score on B then swap; average NLL and ECE (§0.4)."""
+    items = []
+    for r in rows:
+        if "logits" not in r or "labels" not in r:
+            continue
+        items.append({
+            "group": group_of(r.get("record_id"), r.get("view_group")),
+            "type": r["type"], "kb": k_bucket(len(r["labels"])),
+            "logits": torch.tensor(r["logits"]), "idx": r["labels"].index(r["gold"]),
+        })
+    fold_of = [_fold_of(it["group"], folds) for it in items]
+    nlls, eces = [], []
+    for f in range(folds):
+        fit = [it for it, ff in zip(items, fold_of) if ff != f]
+        ev = [it for it, ff in zip(items, fold_of) if ff == f]
+        temps = _fit_temps_from_items(fit, min_items)
+        n, e, _, _ = _score_items(ev, temps)
+        nlls.append(n); eces.append(e)
+    return {"nll_cal_cv": round(sum(nlls) / folds, 4),
+            "ece_cal_cv": round(sum(eces) / folds, 4)}

nomeda/config.py

@@ -0,0 +1,103 @@
+"""NOMEDA-Decision v1 configuration.
+
+All hyper-parameters live here and are loaded from YAML. Defaults match the
+implementation spec.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field, asdict
+from pathlib import Path
+from typing import List, Optional
+
+import yaml
+
+ENCODER_ID = "silma-ai/silma-embedding-matryoshka-v0.1"
+# Pinned revision (verified, see docs/VERIFIED.md)
+ENCODER_REVISION = "011e3b07c9ae2e38f2312a237b489b3b0a640307"
+
+K_BUCKETS = [(2, 2), (3, 5), (6, 10), (11, 10**9)]
+K_BUCKET_NAMES = ["2", "3-5", "6-10", "11+"]
+
+
+def k_bucket(k: int) -> str:
+    for (lo, hi), name in zip(K_BUCKETS, K_BUCKET_NAMES):
+        if lo <= k <= hi:
+            return name
+    raise ValueError(k)
+
+
+@dataclass
+class NomedaConfig:
+    # encoder
+    encoder_id: str = ENCODER_ID
+    encoder_revision: str = ENCODER_REVISION
+    encoder_dim: int = 768
+    encoder_layers: int = 12
+    encoder_max_len: int = 512
+    encoder_vocab: int = 64000
+
+    # decision layer
+    d_model: int = 384
+    n_heads: int = 6
+    ffn_mult: int = 4
+    n_decision_blocks: int = 4
+    dropout: float = 0.1
+    option_state_xattn: bool = True
+
+    # state chunking
+    state_max_len: int = 512
+    state_chunk_stride: int = 384
+    state_max_chunks: int = 8
+
+    # text
+    text_max_len: int = 128
+    n_rank_levels: int = 32
+    q_bucket: int = 128  # questions per length-bucketed sub-batch (padding control)
+
+    # freezing
+    freeze_encoder: bool = True
+    unfreeze_top_layers: int = 0
+    # Optional GRADUAL unfreeze. When non-empty, e.g. [0, 2, 4, 6], the number of
+    # trainable top encoder layers is re-set at the start of each epoch to the
+    # given value (last value repeats if epochs exceed the list). Empty = fixed
+    # `unfreeze_top_layers` for the whole run. Only used as an escalation after
+    # fixed unfreezing fails to earn its place (plan C4.4).
+    unfreeze_schedule: List[int] = field(default_factory=list)
+
+    # loss
+    label_smoothing: float = 0.02
+    w_spherical: float = 0.5
+    w_rps: float = 1.0
+    lambda_kl: float = 0.0
+    kl_en_ar_only: bool = True  # restrict consistency KL to V2 en<->ar pairs
+
+    # training
+    lr_decision: float = 3e-4
+    lr_encoder: float = 2e-5
+    weight_decay: float = 0.01
+    warmup_frac: float = 0.03
+    grad_clip: float = 1.0
+    bf16: bool = True
+    seed: int = 13
+
+    @classmethod
+    def from_yaml(cls, path: str | Path) -> "NomedaConfig":
+        data = yaml.safe_load(Path(path).read_text(encoding="utf-8")) or {}
+        known = {f for f in cls.__dataclass_fields__}
+        unknown = set(data) - known
+        if unknown:
+            raise ValueError(f"unknown config keys: {sorted(unknown)}")
+        return cls(**data)
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+DEFAULT_CONFIG = NomedaConfig()
+
+
+# ---- backward-compatible aliases -----------------------------------------
+# The class was originally published as `NomedConfig` in NOMEDA-Decision
+# v0.1-beta. Keep the old name working so existing code and the published model
+# card do not break; new code should use `NomedaConfig`.
+NomedConfig = NomedaConfig

nomeda/costs.py

@@ -0,0 +1,93 @@
+"""Cost-aware decisions (§Stage 7, step 4a).
+
+Given calibrated probabilities p(y) over the supplied options and a cost matrix
+C[action, y], the Bayes-optimal single-step action is
+
+    a* = argmin_a  Σ_y p(y) · C[a, y]
+
+This is exact when the probabilities are calibrated, needs no training, and is the
+baseline that the RL policy (4b-4d) must beat.
+
+Cost design is an owner product decision: matrices live in ``costs/*.yaml``, are versioned,
+and carry the reason for each number. Nothing here ships a default matrix silently.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List, Optional, Sequence, Tuple
+
+import yaml
+
+
+@dataclass
+class CostMatrix:
+    """Actions x outcomes cost matrix for one question.
+
+    actions: ordered action ids.
+    matrix:  {action: {option_label: cost}}; every action must price every option.
+    """
+    actions: List[str]
+    matrix: Dict[str, Dict[str, float]]
+    version: str = "unknown"
+    notes: str = ""
+
+    def validate(self, option_labels: Sequence[str]) -> None:
+        if not self.actions:
+            raise ValueError("cost matrix has no actions")
+        for a in self.actions:
+            if a not in self.matrix:
+                raise ValueError(f"cost matrix missing action {a!r}")
+            row = self.matrix[a]
+            missing = [lab for lab in option_labels if lab not in row]
+            if missing:
+                raise ValueError(f"action {a!r} does not price options {missing}")
+            for lab in option_labels:
+                c = row[lab]
+                if not isinstance(c, (int, float)) or c != c:
+                    raise ValueError(f"cost for {a!r}/{lab!r} is not a number: {c!r}")
+                if c < 0:
+                    raise ValueError(f"cost for {a!r}/{lab!r} is negative: {c}")
+
+    def as_features(self, option_labels: Sequence[str]) -> List[List[float]]:
+        """Matrix rows in canonical option order (for the policy head)."""
+        return [[float(self.matrix[a][lab]) for lab in option_labels] for a in self.actions]
+
+
+def expected_cost(probs: Sequence[float], costs: Sequence[float]) -> float:
+    """Σ_y p(y) · C[a, y] for one action."""
+    if len(probs) != len(costs):
+        raise ValueError(f"length mismatch: {len(probs)} probs vs {len(costs)} costs")
+    return float(sum(p * c for p, c in zip(probs, costs)))
+
+
+def bayes_action(probs: Sequence[float], labels: Sequence[str],
+                 cost: CostMatrix) -> Tuple[str, float, Dict[str, float]]:
+    """Return (best_action, its expected cost, {action: expected cost})."""
+    cost.validate(labels)
+    ec = {a: expected_cost(probs, [cost.matrix[a][lab] for lab in labels]) for a in cost.actions}
+    best = min(cost.actions, key=lambda a: (ec[a], a))  # lexicographic tie-break
+    return best, ec[best], ec
+
+
+def load_costs(path: str | Path, option_labels: Sequence[str]) -> CostMatrix:
+    """Load and validate a cost matrix from YAML for a given option list."""
+    data = yaml.safe_load(Path(path).read_text(encoding="utf-8")) or {}
+    cm = CostMatrix(actions=list(data.get("actions", [])),
+                    matrix=data.get("matrix", {}),
+                    version=str(data.get("version", "unknown")),
+                    notes=str(data.get("notes", "")))
+    cm.validate(option_labels)
+    return cm
+
+
+def apply_costs(answer: dict, cost: CostMatrix, labels: Sequence[str]) -> dict:
+    """Attach `action`, `expected_cost` and per-action expected costs to a runtime answer."""
+    probs = [float(answer["probabilities"][lab]) for lab in labels]
+    best, best_ec, ec = bayes_action(probs, labels, cost)
+    out = dict(answer)
+    out["action"] = best
+    out["expected_cost"] = round(best_ec, 6)
+    out["expected_cost_by_action"] = {a: round(v, 6) for a, v in ec.items()}
+    out["cost_version"] = cost.version
+    return out

nomeda/data/__init__.py

@@ -0,0 +1,2 @@
+from .loader import iter_records, record_to_example, family_of, resolve_gold, strip_evidence, required_split_paths, FORBIDDEN  # noqa
+from .batching import iter_batches, augment, pair_key  # noqa