qarar 1.0.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (45) hide show
  1. qarar-1.0.2/PKG-INFO +179 -0
  2. qarar-1.0.2/README.md +149 -0
  3. qarar-1.0.2/nomeda/__init__.py +71 -0
  4. qarar-1.0.2/nomeda/blocks.py +112 -0
  5. qarar-1.0.2/nomeda/calibrate.py +165 -0
  6. qarar-1.0.2/nomeda/config.py +103 -0
  7. qarar-1.0.2/nomeda/costs.py +93 -0
  8. qarar-1.0.2/nomeda/data/__init__.py +2 -0
  9. qarar-1.0.2/nomeda/data/batching.py +194 -0
  10. qarar-1.0.2/nomeda/data/family_map.yaml +51 -0
  11. qarar-1.0.2/nomeda/data/loader.py +119 -0
  12. qarar-1.0.2/nomeda/data/variants.py +274 -0
  13. qarar-1.0.2/nomeda/data/weights.py +48 -0
  14. qarar-1.0.2/nomeda/encoder.py +152 -0
  15. qarar-1.0.2/nomeda/features.py +255 -0
  16. qarar-1.0.2/nomeda/loss.py +125 -0
  17. qarar-1.0.2/nomeda/metrics.py +217 -0
  18. qarar-1.0.2/nomeda/model.py +224 -0
  19. qarar-1.0.2/nomeda/partial_cache.py +280 -0
  20. qarar-1.0.2/nomeda/registry.py +170 -0
  21. qarar-1.0.2/nomeda/runtime.py +204 -0
  22. qarar-1.0.2/nomeda/text.py +97 -0
  23. qarar-1.0.2/nomeda/trainer.py +255 -0
  24. qarar-1.0.2/pyproject.toml +47 -0
  25. qarar-1.0.2/qarar/__init__.py +34 -0
  26. qarar-1.0.2/qarar.egg-info/PKG-INFO +179 -0
  27. qarar-1.0.2/qarar.egg-info/SOURCES.txt +43 -0
  28. qarar-1.0.2/qarar.egg-info/dependency_links.txt +1 -0
  29. qarar-1.0.2/qarar.egg-info/requires.txt +19 -0
  30. qarar-1.0.2/qarar.egg-info/top_level.txt +2 -0
  31. qarar-1.0.2/setup.cfg +4 -0
  32. qarar-1.0.2/tests/test_chunking_correctness.py +164 -0
  33. qarar-1.0.2/tests/test_costs.py +147 -0
  34. qarar-1.0.2/tests/test_determinism.py +34 -0
  35. qarar-1.0.2/tests/test_features.py +68 -0
  36. qarar-1.0.2/tests/test_flat_loss.py +41 -0
  37. qarar-1.0.2/tests/test_invariance.py +62 -0
  38. qarar-1.0.2/tests/test_leakage.py +34 -0
  39. qarar-1.0.2/tests/test_metrics.py +48 -0
  40. qarar-1.0.2/tests/test_pooling_parity.py +43 -0
  41. qarar-1.0.2/tests/test_registry.py +57 -0
  42. qarar-1.0.2/tests/test_runtime.py +97 -0
  43. qarar-1.0.2/tests/test_trainer.py +115 -0
  44. qarar-1.0.2/tests/test_variants.py +143 -0
  45. qarar-1.0.2/tests/test_weights.py +70 -0
qarar-1.0.2/PKG-INFO ADDED
@@ -0,0 +1,179 @@
1
+ Metadata-Version: 2.4
2
+ Name: qarar
3
+ Version: 1.0.2
4
+ Summary: QARAR: Arabic-first typed-decision models (calibrated distributions over supplied options).
5
+ Author: NOMEDA
6
+ License-Expression: Apache-2.0
7
+ Project-URL: Homepage, https://huggingface.co/nomeda-lab/nomeda-decision-v0.1-beta
8
+ Keywords: arabic,nlp,typed-decisions,calibration,probability-estimation
9
+ Classifier: Programming Language :: Python :: 3
10
+ Classifier: Natural Language :: Arabic
11
+ Classifier: Natural Language :: English
12
+ Requires-Python: >=3.10
13
+ Description-Content-Type: text/markdown
14
+ Requires-Dist: torch>=2.2
15
+ Requires-Dist: transformers>=4.44
16
+ Requires-Dist: safetensors>=0.4
17
+ Requires-Dist: huggingface_hub>=0.23
18
+ Requires-Dist: numpy>=1.24
19
+ Requires-Dist: PyYAML
20
+ Provides-Extra: data
21
+ Requires-Dist: datasets>=2.18; extra == "data"
22
+ Requires-Dist: pandas>=2.0; extra == "data"
23
+ Requires-Dist: pyarrow>=14; extra == "data"
24
+ Provides-Extra: train
25
+ Requires-Dist: sentence-transformers>=3.0; extra == "train"
26
+ Requires-Dist: regex; extra == "train"
27
+ Requires-Dist: pydantic>=2.5; extra == "train"
28
+ Provides-Extra: dev
29
+ Requires-Dist: pytest; extra == "dev"
30
+
31
+ # NOMEDA-Decision v1
32
+
33
+ Arabic-first **typed-decision** model. Given one *state* it answers **many independent
34
+ questions** by returning a calibrated probability distribution over the **supplied
35
+ options** — it never generates text.
36
+
37
+ ## Install
38
+
39
+ ```bash
40
+ pip install qarar
41
+ ```
42
+
43
+ ```python
44
+ from qarar import load # published alias; `from nomeda import load` also works
45
+ m = load("nomeda-lab/qarar")
46
+ ```
47
+
48
+ Question types share one scoring path:
49
+
50
+ | type | meaning | options |
51
+ |---|---|---|
52
+ | `choice` | pick one named option | K named options |
53
+ | `score` | rank/level a target | K ordered options (+ `rank_emb`) |
54
+ | `noul` | yes / no / unknown | 2 options (`false`, `true`) |
55
+
56
+ - Encoder: `silma-ai/silma-embedding-matryoshka-v0.1` (Arabetv02 backbone, D=768, mean
57
+ pooling), loaded via `AutoModel`/`AutoTokenizer` (no `trust_remote_code`).
58
+ - Non-interference is **exact**: each question cross-attends only to the shared state,
59
+ never to another question (block-diagonal attention).
60
+ - Encoder calls per request ≤ 2: state once, all question/option texts in one batched pass.
61
+
62
+ See `docs/MODEL_REPORT.md` for the architecture, measured parameter budget, latency and
63
+ readiness. Diagram: `docs/nomeda_decision_architecture.drawio`.
64
+
65
+ ## Layout
66
+
67
+ ```
68
+ nomeda/ library
69
+ config.py NomedaConfig
70
+ encoder.py SilmaEncoder (mean pooling, chunking, LRU cache)
71
+ blocks.py DecisionBlock / OptionStateBlock / AttnPool / CompatibilityScorer
72
+ model.py NomedaModel + param_report()
73
+ runtime.py NomedaRuntime.predict / predict_batch, validate_questions
74
+ loss.py proper-scoring losses
75
+ calibrate.py temperature scaling, abstention thresholds, ECE
76
+ data/ loader, batching/augmentation, family_map.yaml
77
+ scripts/ audit_data, train, calibrate, evaluate, export, bench_latency
78
+ tests/ invariance, runtime, leakage, pooling-parity, determinism
79
+ configs/ default.yaml
80
+ docs/ MODEL_REPORT, DATA_AUDIT, VERIFIED, OPEN_QUESTIONS, RESULTS
81
+ ```
82
+
83
+ ## Install
84
+
85
+ ```bash
86
+ pip install nomeda # from PyPI (v1.0.1+)
87
+ # or, for development, from the repository root:
88
+ pip install -e .
89
+ # with the data-prep and training extras:
90
+ pip install -e ".[data,train,dev]"
91
+ ```
92
+
93
+ `pip install nomeda` (or `pip install -e .`) makes `import nomeda` work from any
94
+ directory.
95
+
96
+ ## Usage
97
+
98
+ **Recommended — load the published weights (calibration included):**
99
+
100
+ ```python
101
+ from nomeda import load
102
+
103
+ nm = load("nomeda-lab/nomeda-decision-v0.1-beta") # or a local export directory
104
+ out = nm.predict(
105
+ state="أنا اتخصم مني مرتين على فاتورة مارس، عايز أرجّع الفلوس النهاردة.",
106
+ questions={
107
+ "intent": {"type": "choice", "instructions": "ماذا يريد العميل؟",
108
+ "criteria": {"refund": "استرداد المال",
109
+ "technical_help": "مشكلة تقنية",
110
+ "cancellation": "إلغاء الخدمة"}},
111
+ "frustration": {"type": "score", "instructions": "ما مدى انزعاج العميل؟",
112
+ "criteria": ["هادئ", "قلق", "منزعج", "غاضب جدًا"]},
113
+ "refund_requested": {"type": "noul", "instructions": "هل يطلب العميل استرداد أمواله؟"},
114
+ },
115
+ )
116
+
117
+ out["answers"]["intent"]["choice"] # -> "refund"
118
+ out["answers"]["intent"]["probabilities"] # -> {"refund": 0.555, ...}
119
+ out["answers"]["refund_requested"]["noul"] # -> 0.608
120
+ ```
121
+
122
+ `load()` fetches `temperatures.json` and `thresholds.json` alongside the weights,
123
+ so predictions are calibrated and abstention works. **This is what you want
124
+ almost always.**
125
+
126
+ **Manual construction** (when you have the weights and calibration already and
127
+ want full control):
128
+
129
+ ```python
130
+ import json
131
+ from nomeda import NomedaConfig, NomedaModel
132
+ from nomeda.runtime import NomedaRuntime
133
+
134
+ model = NomedaModel(NomedaConfig(), device="cpu") # or device="cuda"
135
+
136
+ # temperatures/thresholds are REAL files, not placeholders. Load them:
137
+ temperatures = json.load(open("temperatures.json")) # {type: {bucket: T}}
138
+ thresholds = json.load(open("thresholds.json")) # {type: p_min}
139
+
140
+ rt = NomedaRuntime(model, temperatures=temperatures,
141
+ thresholds=thresholds, device="cpu")
142
+ rt.predict(state, questions)
143
+ ```
144
+
145
+ Without calibration pass `temperatures={}, thresholds={}` (identity temperature,
146
+ no abstention) — never `...`, which raises a `TypeError`.
147
+
148
+ > **Naming.** The classes were `Nomed*` in the v0.1-beta release. `NomedConfig`,
149
+ > `NomedModel`, `NomedRuntime` and `Nomed` remain available as aliases, so existing
150
+ > code keeps working; new code should use the `Nomeda*` names.
151
+
152
+ ## Pipeline
153
+
154
+ ```bash
155
+ DATA=/path/to/data/final_training # train.jsonl, validation.jsonl, test.jsonl
156
+ python scripts/train.py --data-dir $DATA --out checkpoints/run
157
+ python scripts/calibrate.py --data-dir $DATA --checkpoint checkpoints/run/checkpoint.pt --out artifacts
158
+ python scripts/evaluate.py --data-dir $DATA --checkpoint checkpoints/run/checkpoint.pt --artifacts artifacts --out docs/RESULTS.md
159
+ python scripts/export.py --checkpoint checkpoints/run/checkpoint.pt --artifacts artifacts --out dist/nomeda-decision
160
+ ```
161
+
162
+ ## Tests
163
+
164
+ ```bash
165
+ PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest tests -m "not slow" # fast suite
166
+ PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest tests -m "slow" # real-encoder parity
167
+ ```
168
+
169
+ ## Kaggle
170
+
171
+ `kaggle/nomeda_decision_kaggle.ipynb` runs the whole frozen-encoder pipeline (feature
172
+ cache → baselines → training → calibration → test → export/handoff) on a Kaggle GPU.
173
+ Upload the repo as a zip dataset and the `*.jsonl` data as another dataset, attach both,
174
+ enable GPU + Internet, and run the cells top to bottom. See the notebook's first cell.
175
+
176
+ ## License
177
+
178
+ Code: see `LICENSES.yaml` / repository. Encoder `silma-embedding-matryoshka-v0.1` is
179
+ Apache-2.0 (model card; no standalone LICENSE file upstream).
qarar-1.0.2/README.md ADDED
@@ -0,0 +1,149 @@
1
+ # NOMEDA-Decision v1
2
+
3
+ Arabic-first **typed-decision** model. Given one *state* it answers **many independent
4
+ questions** by returning a calibrated probability distribution over the **supplied
5
+ options** — it never generates text.
6
+
7
+ ## Install
8
+
9
+ ```bash
10
+ pip install qarar
11
+ ```
12
+
13
+ ```python
14
+ from qarar import load # published alias; `from nomeda import load` also works
15
+ m = load("nomeda-lab/qarar")
16
+ ```
17
+
18
+ Question types share one scoring path:
19
+
20
+ | type | meaning | options |
21
+ |---|---|---|
22
+ | `choice` | pick one named option | K named options |
23
+ | `score` | rank/level a target | K ordered options (+ `rank_emb`) |
24
+ | `noul` | yes / no / unknown | 2 options (`false`, `true`) |
25
+
26
+ - Encoder: `silma-ai/silma-embedding-matryoshka-v0.1` (Arabetv02 backbone, D=768, mean
27
+ pooling), loaded via `AutoModel`/`AutoTokenizer` (no `trust_remote_code`).
28
+ - Non-interference is **exact**: each question cross-attends only to the shared state,
29
+ never to another question (block-diagonal attention).
30
+ - Encoder calls per request ≤ 2: state once, all question/option texts in one batched pass.
31
+
32
+ See `docs/MODEL_REPORT.md` for the architecture, measured parameter budget, latency and
33
+ readiness. Diagram: `docs/nomeda_decision_architecture.drawio`.
34
+
35
+ ## Layout
36
+
37
+ ```
38
+ nomeda/ library
39
+ config.py NomedaConfig
40
+ encoder.py SilmaEncoder (mean pooling, chunking, LRU cache)
41
+ blocks.py DecisionBlock / OptionStateBlock / AttnPool / CompatibilityScorer
42
+ model.py NomedaModel + param_report()
43
+ runtime.py NomedaRuntime.predict / predict_batch, validate_questions
44
+ loss.py proper-scoring losses
45
+ calibrate.py temperature scaling, abstention thresholds, ECE
46
+ data/ loader, batching/augmentation, family_map.yaml
47
+ scripts/ audit_data, train, calibrate, evaluate, export, bench_latency
48
+ tests/ invariance, runtime, leakage, pooling-parity, determinism
49
+ configs/ default.yaml
50
+ docs/ MODEL_REPORT, DATA_AUDIT, VERIFIED, OPEN_QUESTIONS, RESULTS
51
+ ```
52
+
53
+ ## Install
54
+
55
+ ```bash
56
+ pip install nomeda # from PyPI (v1.0.1+)
57
+ # or, for development, from the repository root:
58
+ pip install -e .
59
+ # with the data-prep and training extras:
60
+ pip install -e ".[data,train,dev]"
61
+ ```
62
+
63
+ `pip install nomeda` (or `pip install -e .`) makes `import nomeda` work from any
64
+ directory.
65
+
66
+ ## Usage
67
+
68
+ **Recommended — load the published weights (calibration included):**
69
+
70
+ ```python
71
+ from nomeda import load
72
+
73
+ nm = load("nomeda-lab/nomeda-decision-v0.1-beta") # or a local export directory
74
+ out = nm.predict(
75
+ state="أنا اتخصم مني مرتين على فاتورة مارس، عايز أرجّع الفلوس النهاردة.",
76
+ questions={
77
+ "intent": {"type": "choice", "instructions": "ماذا يريد العميل؟",
78
+ "criteria": {"refund": "استرداد المال",
79
+ "technical_help": "مشكلة تقنية",
80
+ "cancellation": "إلغاء الخدمة"}},
81
+ "frustration": {"type": "score", "instructions": "ما مدى انزعاج العميل؟",
82
+ "criteria": ["هادئ", "قلق", "منزعج", "غاضب جدًا"]},
83
+ "refund_requested": {"type": "noul", "instructions": "هل يطلب العميل استرداد أمواله؟"},
84
+ },
85
+ )
86
+
87
+ out["answers"]["intent"]["choice"] # -> "refund"
88
+ out["answers"]["intent"]["probabilities"] # -> {"refund": 0.555, ...}
89
+ out["answers"]["refund_requested"]["noul"] # -> 0.608
90
+ ```
91
+
92
+ `load()` fetches `temperatures.json` and `thresholds.json` alongside the weights,
93
+ so predictions are calibrated and abstention works. **This is what you want
94
+ almost always.**
95
+
96
+ **Manual construction** (when you have the weights and calibration already and
97
+ want full control):
98
+
99
+ ```python
100
+ import json
101
+ from nomeda import NomedaConfig, NomedaModel
102
+ from nomeda.runtime import NomedaRuntime
103
+
104
+ model = NomedaModel(NomedaConfig(), device="cpu") # or device="cuda"
105
+
106
+ # temperatures/thresholds are REAL files, not placeholders. Load them:
107
+ temperatures = json.load(open("temperatures.json")) # {type: {bucket: T}}
108
+ thresholds = json.load(open("thresholds.json")) # {type: p_min}
109
+
110
+ rt = NomedaRuntime(model, temperatures=temperatures,
111
+ thresholds=thresholds, device="cpu")
112
+ rt.predict(state, questions)
113
+ ```
114
+
115
+ Without calibration pass `temperatures={}, thresholds={}` (identity temperature,
116
+ no abstention) — never `...`, which raises a `TypeError`.
117
+
118
+ > **Naming.** The classes were `Nomed*` in the v0.1-beta release. `NomedConfig`,
119
+ > `NomedModel`, `NomedRuntime` and `Nomed` remain available as aliases, so existing
120
+ > code keeps working; new code should use the `Nomeda*` names.
121
+
122
+ ## Pipeline
123
+
124
+ ```bash
125
+ DATA=/path/to/data/final_training # train.jsonl, validation.jsonl, test.jsonl
126
+ python scripts/train.py --data-dir $DATA --out checkpoints/run
127
+ python scripts/calibrate.py --data-dir $DATA --checkpoint checkpoints/run/checkpoint.pt --out artifacts
128
+ python scripts/evaluate.py --data-dir $DATA --checkpoint checkpoints/run/checkpoint.pt --artifacts artifacts --out docs/RESULTS.md
129
+ python scripts/export.py --checkpoint checkpoints/run/checkpoint.pt --artifacts artifacts --out dist/nomeda-decision
130
+ ```
131
+
132
+ ## Tests
133
+
134
+ ```bash
135
+ PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest tests -m "not slow" # fast suite
136
+ PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest tests -m "slow" # real-encoder parity
137
+ ```
138
+
139
+ ## Kaggle
140
+
141
+ `kaggle/nomeda_decision_kaggle.ipynb` runs the whole frozen-encoder pipeline (feature
142
+ cache → baselines → training → calibration → test → export/handoff) on a Kaggle GPU.
143
+ Upload the repo as a zip dataset and the `*.jsonl` data as another dataset, attach both,
144
+ enable GPU + Internet, and run the cells top to bottom. See the notebook's first cell.
145
+
146
+ ## License
147
+
148
+ Code: see `LICENSES.yaml` / repository. Encoder `silma-embedding-matryoshka-v0.1` is
149
+ Apache-2.0 (model card; no standalone LICENSE file upstream).
@@ -0,0 +1,71 @@
1
+ """NOMEDA-Decision v1: one Arabic state -> many independent typed decisions."""
2
+ from __future__ import annotations
3
+
4
+ import json
5
+ from pathlib import Path
6
+
7
+ from .config import (NomedaConfig, NomedConfig, DEFAULT_CONFIG, k_bucket,
8
+ K_BUCKETS, K_BUCKET_NAMES)
9
+ from .model import NomedaModel, NomedModel
10
+ from .runtime import NomedaRuntime, NomedRuntime, validate_questions
11
+
12
+ __version__ = "1.0.0"
13
+ __all__ = ["load", "Nomeda", "Nomed",
14
+ "NomedaConfig", "NomedaModel", "NomedaRuntime",
15
+ "NomedConfig", "NomedModel", "NomedRuntime",
16
+ "validate_questions", "k_bucket"]
17
+
18
+
19
+ class Nomeda:
20
+ """Loaded model + calibration; ``predict`` mirrors the Laya output shape."""
21
+
22
+ def __init__(self, model: NomedaModel, temperatures: dict | None = None,
23
+ thresholds: dict | None = None, device: str = "cpu"):
24
+ self.model = model
25
+ self.encoder_id = model.cfg.encoder_id
26
+ self.encoder_revision = model.cfg.encoder_revision
27
+ self.runtime = NomedaRuntime(model, temperatures or {}, thresholds or {}, device=device)
28
+
29
+ def predict(self, state, questions):
30
+ return self.runtime.predict(state, questions)
31
+
32
+ def predict_batch(self, states, questions):
33
+ return self.runtime.predict_batch(states, questions)
34
+
35
+ def param_report(self) -> dict:
36
+ return self.model.param_report()
37
+
38
+
39
+ def load(path_or_hub_id: str, device: str = "auto", tiny: bool = False) -> Nomeda:
40
+ """Load a NOMEDA-Decision model from a local export directory or a Hub id."""
41
+ if device == "auto":
42
+ import torch
43
+ device = "cuda" if torch.cuda.is_available() else "cpu"
44
+ if Path(path_or_hub_id).is_dir():
45
+ root = Path(path_or_hub_id)
46
+ raw = json.loads((root / "config.json").read_text(encoding="utf-8"))
47
+ cfg = NomedaConfig(**{k: v for k, v in raw.items() if k in NomedaConfig.__dataclass_fields__})
48
+ enc_dir = root / "encoder"
49
+ model = NomedaModel(cfg, device=device, tiny=tiny)
50
+ if enc_dir.exists():
51
+ from transformers import AutoModel, AutoTokenizer
52
+ model.encoder.tok = AutoTokenizer.from_pretrained(enc_dir)
53
+ model.encoder.model = AutoModel.from_pretrained(enc_dir).to(device)
54
+ head = root / "decision_head.safetensors"
55
+ if head.exists():
56
+ from safetensors.torch import load_file
57
+ state = load_file(str(head))
58
+ model.load_state_dict(state, strict=False)
59
+ temps = json.loads((root / "temperatures.json").read_text(encoding="utf-8")) if (root / "temperatures.json").exists() else {}
60
+ thr = json.loads((root / "thresholds.json").read_text(encoding="utf-8")) if (root / "thresholds.json").exists() else {}
61
+ model.eval()
62
+ model.to(device)
63
+ return Nomeda(model, temps, thr, device=device)
64
+ # Hub id
65
+ from huggingface_hub import snapshot_download
66
+ d = snapshot_download(repo_id=path_or_hub_id)
67
+ return load(d, device=device, tiny=tiny)
68
+
69
+
70
+ # ---- backward-compatible alias (published as `Nomed` in v0.1-beta) --------
71
+ Nomed = Nomeda
@@ -0,0 +1,112 @@
1
+ """Decision blocks, option-state block and attention pooling (§3.4–3.6)."""
2
+ from __future__ import annotations
3
+
4
+ import torch
5
+ import torch.nn as nn
6
+ import torch.nn.functional as F
7
+
8
+
9
+ class MultiheadAttention(nn.Module):
10
+ def __init__(self, d: int, n_heads: int, dropout: float = 0.0):
11
+ super().__init__()
12
+ assert d % n_heads == 0
13
+ self.d, self.h = d, n_heads
14
+ self.dh = d // n_heads
15
+ self.q = nn.Linear(d, d)
16
+ self.k = nn.Linear(d, d)
17
+ self.v = nn.Linear(d, d)
18
+ self.o = nn.Linear(d, d)
19
+ self.dropout = dropout
20
+
21
+ def forward(self, q, k, v, q_mask=None, k_mask=None):
22
+ B, Lq, _ = q.shape
23
+ Lk = k.shape[1]
24
+ Q = self.q(q).view(B, Lq, self.h, self.dh).transpose(1, 2)
25
+ K = self.k(k).view(B, Lk, self.h, self.dh).transpose(1, 2)
26
+ V = self.v(v).view(B, Lk, self.h, self.dh).transpose(1, 2)
27
+ attn_mask = None
28
+ if k_mask is not None:
29
+ attn_mask = torch.zeros(B, 1, 1, Lk, device=q.device, dtype=Q.dtype)
30
+ attn_mask = attn_mask.masked_fill(~k_mask[:, None, None, :].bool(), -1e4)
31
+ out = F.scaled_dot_product_attention(
32
+ Q, K, V, attn_mask=attn_mask,
33
+ dropout_p=self.dropout if self.training else 0.0)
34
+ out = out.transpose(1, 2).reshape(B, Lq, self.d)
35
+ return self.o(out)
36
+
37
+
38
+ class FFN(nn.Module):
39
+ def __init__(self, d: int, mult: int, dropout: float):
40
+ super().__init__()
41
+ self.net = nn.Sequential(nn.Linear(d, d * mult), nn.GELU(), nn.Dropout(dropout), nn.Linear(d * mult, d))
42
+
43
+ def forward(self, x):
44
+ return self.net(x)
45
+
46
+
47
+ class DecisionBlock(nn.Module):
48
+ """Self-attn within a question -> cross-attn to its own state -> FFN (pre-norm)."""
49
+
50
+ def __init__(self, d, n_heads, mult, dropout):
51
+ super().__init__()
52
+ self.ln1, self.ln2, self.ln3 = nn.LayerNorm(d), nn.LayerNorm(d), nn.LayerNorm(d)
53
+ self.self_attn = MultiheadAttention(d, n_heads, dropout)
54
+ self.cross_attn = MultiheadAttention(d, n_heads, dropout)
55
+ self.ffn = FFN(d, mult, dropout)
56
+
57
+ def forward(self, Q, q_mask, S_pad, S_mask, state_idx):
58
+ Q = Q + self.self_attn(self.ln1(Q), self.ln1(Q), self.ln1(Q), q_mask=q_mask, k_mask=q_mask)
59
+ # cross-attention grouped by state: gather each question's own state (vectorised).
60
+ # Equivalent to a per-state loop but without one kernel launch + sync per state.
61
+ x = self.ln2(Q)
62
+ k = S_pad.index_select(0, state_idx)
63
+ kmask = S_mask.index_select(0, state_idx)
64
+ Q = Q + self.cross_attn(x, k, k, k_mask=kmask)
65
+ Q = Q + self.ffn(self.ln3(Q))
66
+ return Q
67
+
68
+
69
+ class OptionStateBlock(nn.Module):
70
+ def __init__(self, d, n_heads, mult, dropout):
71
+ super().__init__()
72
+ self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d)
73
+ self.cross_attn = MultiheadAttention(d, n_heads, dropout)
74
+ self.ffn = FFN(d, mult, dropout)
75
+
76
+ def forward(self, E, S_pad, S_mask, state_idx):
77
+ x = self.ln1(E).unsqueeze(1) # [Nopt,1,d]
78
+ k = S_pad.index_select(0, state_idx)
79
+ kmask = S_mask.index_select(0, state_idx)
80
+ out = self.cross_attn(x, k, k, k_mask=kmask)
81
+ E = E + out.squeeze(1)
82
+ E = E + self.ffn(self.ln2(E))
83
+ return E
84
+
85
+
86
+ class AttnPool(nn.Module):
87
+ def __init__(self, d):
88
+ super().__init__()
89
+ self.query = nn.Parameter(torch.randn(d) * 0.02)
90
+ self.ln = nn.LayerNorm(d)
91
+
92
+ def forward(self, Q, mask):
93
+ attn = (Q @ self.query) / (Q.shape[-1] ** 0.5) # [B,L]
94
+ attn = attn.masked_fill(~mask.bool(), -1e4)
95
+ p = torch.softmax(attn, dim=-1).unsqueeze(1)
96
+ pooled = torch.bmm(p, Q).squeeze(1)
97
+ return self.ln(pooled)
98
+
99
+
100
+ class CompatibilityScorer(nn.Module):
101
+ def __init__(self, d, dropout):
102
+ super().__init__()
103
+ self.ln = nn.LayerNorm(4 * d)
104
+ self.w1 = nn.Linear(4 * d, d)
105
+ self.w2 = nn.Linear(d, 1)
106
+ self.dropout = nn.Dropout(dropout)
107
+
108
+ def forward(self, q, E):
109
+ h = torch.cat([q, E, q * E, (q - E).abs()], dim=-1)
110
+ h = self.ln(h)
111
+ h = self.dropout(F.gelu(self.w1(h)))
112
+ return self.w2(h).squeeze(-1)
@@ -0,0 +1,165 @@
1
+ """Calibration: temperatures per (type, K-bucket) and abstention thresholds (§6)."""
2
+ from __future__ import annotations
3
+
4
+ import hashlib
5
+ import json
6
+ import math
7
+ from collections import defaultdict
8
+ from typing import Dict, List, Optional, Tuple
9
+
10
+ import torch
11
+ import torch.nn.functional as F
12
+
13
+ from .config import k_bucket
14
+
15
+
16
+ def ece(probs: List[float], correct: List[int], n_bins: int = 15) -> float:
17
+ """Expected calibration error, confidence = max(p), equal-width bins (Laya's definition)."""
18
+ if not probs:
19
+ return 0.0
20
+ total = len(probs)
21
+ out = 0.0
22
+ for b in range(n_bins):
23
+ lo, hi = b / n_bins, (b + 1) / n_bins
24
+ idx = [i for i, p in enumerate(probs) if (p > lo or (b == 0 and p >= lo)) and p <= hi]
25
+ if not idx:
26
+ continue
27
+ conf = sum(probs[i] for i in idx) / len(idx)
28
+ acc = sum(correct[i] for i in idx) / len(idx)
29
+ out += (len(idx) / total) * abs(conf - acc)
30
+ return out
31
+
32
+
33
+ def fit_temperature(logits_list, targets_list, clamp: Tuple[float, float] = (0.5, 5.0)) -> float:
34
+ """Minimise NLL over a single temperature; logits_list holds per-item 1-D tensors
35
+ (K may vary between items)."""
36
+ logT = torch.zeros(1, requires_grad=True)
37
+ opt = torch.optim.LBFGS([logT], lr=0.1, max_iter=50)
38
+ tg = [int(t) for t in targets_list]
39
+
40
+ def closure():
41
+ opt.zero_grad()
42
+ T = torch.exp(logT).clamp(*clamp)
43
+ loss = 0.0
44
+ for l, t in zip(logits_list, tg):
45
+ loss = loss + F.cross_entropy((l / T).unsqueeze(0).float(), torch.tensor([t]))
46
+ loss = loss / max(1, len(logits_list))
47
+ loss.backward()
48
+ return loss
49
+
50
+ opt.step(closure)
51
+ with torch.no_grad():
52
+ return float(torch.exp(logT).clamp(*clamp))
53
+
54
+
55
+ def fit_temperatures(buckets: Dict[Tuple[str, str], Tuple[list, list]],
56
+ min_items: int = 200) -> dict:
57
+ """buckets: {(type, K-bucket): ([logits_1d...], [targets...])} -> temperatures.json dict."""
58
+ out: Dict[str, Dict[str, float]] = {}
59
+ per_type: Dict[str, Tuple[list, list]] = {}
60
+ for (qtype, kb), (logits, targets) in buckets.items():
61
+ per_type.setdefault(qtype, ([], []))
62
+ per_type[qtype][0].extend(logits)
63
+ per_type[qtype][1].extend(targets)
64
+ if len(logits) < min_items:
65
+ continue
66
+ out.setdefault(qtype, {})[kb] = fit_temperature(logits, targets)
67
+ for qtype, (ls, ts) in per_type.items():
68
+ out.setdefault(qtype, {})["default"] = fit_temperature(ls, ts)
69
+ return out
70
+
71
+
72
+ def fit_abstention(conf: List[float], correct: List[int], target_acc: float = 0.95) -> float:
73
+ """Smallest threshold so answered questions reach accuracy >= target_acc."""
74
+ pairs = sorted(zip(conf, correct), reverse=True)
75
+ if not pairs:
76
+ return 0.0
77
+ best_t = 1.0
78
+ n_ok = 0
79
+ for i, (c, ok) in enumerate(pairs):
80
+ n_ok += ok
81
+ acc = n_ok / (i + 1)
82
+ if acc >= target_acc:
83
+ best_t = c
84
+ return float(best_t)
85
+
86
+
87
+ def coverage(conf: List[float], threshold: float) -> float:
88
+ if not conf:
89
+ return 0.0
90
+ return sum(1 for c in conf if c >= threshold) / len(conf)
91
+
92
+
93
+ def save(path: str, obj: dict) -> None:
94
+ with open(path, "w", encoding="utf-8") as f:
95
+ json.dump(obj, f, ensure_ascii=False, indent=2)
96
+
97
+
98
+ # ---- fair calibrated comparison (cal_cv, §0.4) ---------------------------
99
+ def group_of(record_id, view_group) -> str:
100
+ """Fold-group id: view_group and V2 en↔ar pairs stay together."""
101
+ if view_group:
102
+ return f"vg:{view_group}"
103
+ rid = str(record_id)
104
+ if ":en:" in rid or ":ar:" in rid:
105
+ return "pair:" + rid.replace(":en:", ":*:").replace(":ar:", ":*:")
106
+ return f"solo:{rid}"
107
+
108
+
109
+ def _fold_of(group_id: str, folds: int) -> int:
110
+ h = int(hashlib.sha1(str(group_id).encode()).hexdigest()[:8], 16)
111
+ return h % folds
112
+
113
+
114
+ def _fit_temps_from_items(items: List[dict], min_items: int,
115
+ clamp: Tuple[float, float] = (0.5, 5.0)) -> dict:
116
+ buckets: Dict[Tuple[str, str], Tuple[list, list]] = defaultdict(lambda: ([], []))
117
+ per_type: Dict[str, Tuple[list, list]] = defaultdict(lambda: ([], []))
118
+ for it in items:
119
+ kb = it["kb"]
120
+ buckets[(it["type"], kb)][0].append(it["logits"])
121
+ buckets[(it["type"], kb)][1].append(it["idx"])
122
+ per_type[it["type"]][0].append(it["logits"])
123
+ per_type[it["type"]][1].append(it["idx"])
124
+ out: Dict[str, dict] = {}
125
+ for (qt, kb), (L, T) in buckets.items():
126
+ if len(L) >= min_items:
127
+ out.setdefault(qt, {})[kb] = fit_temperature(L, T, clamp)
128
+ for qt, (L, T) in per_type.items():
129
+ out.setdefault(qt, {})["default"] = fit_temperature(L, T, clamp)
130
+ return out
131
+
132
+
133
+ def _score_items(items: List[dict], temps: dict) -> Tuple[float, float, List[float], List[int]]:
134
+ nll = 0.0
135
+ confs, oks = [], []
136
+ for it in items:
137
+ T = float(temps.get(it["type"], {}).get(it["kb"], temps.get(it["type"], {}).get("default", 1.0)))
138
+ p = torch.softmax(it["logits"] / T, dim=-1)
139
+ nll += -math.log(max(float(p[it["idx"]]), 1e-12))
140
+ confs.append(float(p.max()))
141
+ oks.append(int(int(torch.argmax(p)) == it["idx"]))
142
+ return nll / max(1, len(items)), ece(confs, oks), confs, oks
143
+
144
+
145
+ def cal_cv_from_rows(rows: List[dict], folds: int = 2, min_items: int = 100) -> dict:
146
+ """Fit on fold A / score on B then swap; average NLL and ECE (§0.4)."""
147
+ items = []
148
+ for r in rows:
149
+ if "logits" not in r or "labels" not in r:
150
+ continue
151
+ items.append({
152
+ "group": group_of(r.get("record_id"), r.get("view_group")),
153
+ "type": r["type"], "kb": k_bucket(len(r["labels"])),
154
+ "logits": torch.tensor(r["logits"]), "idx": r["labels"].index(r["gold"]),
155
+ })
156
+ fold_of = [_fold_of(it["group"], folds) for it in items]
157
+ nlls, eces = [], []
158
+ for f in range(folds):
159
+ fit = [it for it, ff in zip(items, fold_of) if ff != f]
160
+ ev = [it for it, ff in zip(items, fold_of) if ff == f]
161
+ temps = _fit_temps_from_items(fit, min_items)
162
+ n, e, _, _ = _score_items(ev, temps)
163
+ nlls.append(n); eces.append(e)
164
+ return {"nll_cal_cv": round(sum(nlls) / folds, 4),
165
+ "ece_cal_cv": round(sum(eces) / folds, 4)}