Update jeba-multi adapter, choice head (r=32), and per-language temperature (v0.2.0)

README.md

@@ -62,19 +62,22 @@ split (`training/fit_calibration.py`). Configs and seed live under `training/con
 Reported by `training/evaluate.py` and rendered by `benchmarks/report.py` (accuracy, ECE, p50/p95
 latency per primitive and language).
 
-**Full-scale run (single RTX 3060 12GB):** 6,000 train / 1,500 eval deterministic synthetic
-records (localized cue terms per language, a learnable `other` team with rich descriptions, and
-distractor clauses), LoRA (r=16) plus a dedicated low-rank `choice` head (near-identity init),
-3 epochs, batch 16, bf16 + gradient checkpointing.
+**Full-scale run (single RTX 3060 12GB):** 9,000 English / 18,000 multilingual train / 1,500 eval
+deterministic synthetic records (fully localized per language, a learnable `other` team with rich
+descriptions, per-record RNG, one-in-six distractor clauses), LoRA (r=16) plus a dedicated
+low-rank `choice` head (r=32, near-identity init), 4 epochs, batch 16, bf16 + gradient
+checkpointing.
 
 | Checkpoint | Accuracy | ECE (calibrated) | p50 (ms) |
 | --- | --- | --- | --- |
-| English (ModernBERT-large + LoRA + choice head) | 0.721 | 0.052 | 18.4 |
-| Multilingual (mmBERT-base + LoRA + choice head) | 0.609 | 0.086 | 12.6 |
-
-Per primitive (English): `choice` 0.782, `noul` 0.718, `score` 0.664; (multilingual): `noul`
-0.728, `score` 0.700, `choice` 0.398. The dedicated `choice` head lifted English `choice` from
-~0.25 (chance) to 0.78; multilingual `choice` and `score` calibration remain the next targets.
+| English (ModernBERT-large + LoRA + choice head) | 0.781 | 0.077 | 22.9 |
+| Multilingual (mmBERT-base + LoRA + choice head) | 0.711 | 0.073 | 12.4 |
+
+Per primitive (English): `choice` 0.708, `noul` 0.744, `score` 0.892; (multilingual): `choice`
+0.734, `noul` 0.716, `score` 0.682. The localized, per-record-RNG data (B-1) lifted multilingual
+`choice` from 0.40 to 0.73 and English overall from 0.72 to 0.78. Per-language calibration for a
+few multilingual languages (`es`, `nl`, `de`) and multilingual `score` remain the next targets.
+The CUDA-graph fast path (`JEBA_FAST=1`) gives a 2.7× p50 speedup with 0 top-label flips.
 Full tables and environment are in
 [`benchmarks/report.md`](https://github.com/munod/jeba/blob/main/benchmarks/report.md).
 

adapter_config.json

@@ -35,8 +35,8 @@
   "rank_pattern": {},
   "revision": null,
   "target_modules": [
-    "Wqkv",
     "Wo",
+    "Wqkv",
     "Wi"
   ],
   "target_parameters": null,

adapter_model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:572a2d107f3656ffeeb6e90783d91e1f2f885b1bb658001b89ea50e036dbbf8c
+oid sha256:6f529ee66b8d859121e075662702384dce21aaddc14d1fdf288b4383c479ae19
 size 13537472

finetune_config.json

@@ -1,17 +1,17 @@
 {
-  "batch_size": 16,
+  "batch_size": 8,
   "bf16": true,
   "choice_init_std": 0.01,
-  "choice_rank": 16,
+  "choice_rank": 32,
   "data_path": "data/train_multi.jsonl",
-  "epochs": 3,
+  "epochs": 4,
   "grad_accum": 4,
   "gradient_checkpointing": true,
   "learning_rate": 0.0001,
   "lora_alpha": 32,
   "lora_dropout": 0.05,
   "lora_rank": 16,
-  "max_len": 256,
+  "max_len": 1024,
   "max_records": null,
   "model_id": "jhu-clsp/mmBERT-base",
   "out_dir": "checkpoints/multi",

temperature_calibration.json

@@ -2,20 +2,134 @@
   "bins": 10,
   "per_primitive": {
     "choice": {
-      "ece_after": 0.037778,
-      "ece_before": 0.141244,
+      "by_language": {
+        "de": {
+          "ece_after": 0.249286,
+          "ece_before": 0.335616,
+          "n": 83,
+          "temperature": 3.0
+        },
+        "es": {
+          "ece_after": 0.295552,
+          "ece_before": 0.395429,
+          "n": 84,
+          "temperature": 3.0
+        },
+        "fr": {
+          "ece_after": 0.031722,
+          "ece_before": 0.031722,
+          "n": 83,
+          "temperature": 1.0
+        },
+        "it": {
+          "ece_after": 0.000804,
+          "ece_before": 0.013927,
+          "n": 83,
+          "temperature": 1.25
+        },
+        "nl": {
+          "ece_after": 0.154009,
+          "ece_before": 0.22301,
+          "n": 83,
+          "temperature": 2.0
+        },
+        "pt": {
+          "ece_after": 0.059923,
+          "ece_before": 0.083558,
+          "n": 84,
+          "temperature": 0.5
+        }
+      },
+      "ece_after": 0.118556,
+      "ece_before": 0.153826,
       "n": 500,
-      "temperature": 4.0
+      "temperature": 2.0
     },
     "noul": {
-      "ece_after": 0.042006,
-      "ece_before": 0.179002,
+      "by_language": {
+        "de": {
+          "ece_after": 0.012124,
+          "ece_before": 0.068245,
+          "n": 83,
+          "temperature": 3.0
+        },
+        "es": {
+          "ece_after": 0.087329,
+          "ece_before": 0.224015,
+          "n": 84,
+          "temperature": 10.0
+        },
+        "fr": {
+          "ece_after": 0.006848,
+          "ece_before": 0.310798,
+          "n": 83,
+          "temperature": 0.1
+        },
+        "it": {
+          "ece_after": 0.040417,
+          "ece_before": 0.259646,
+          "n": 83,
+          "temperature": 0.25
+        },
+        "nl": {
+          "ece_after": 0.078314,
+          "ece_before": 0.218552,
+          "n": 83,
+          "temperature": 10.0
+        },
+        "pt": {
+          "ece_after": 0.006405,
+          "ece_before": 0.29764,
+          "n": 84,
+          "temperature": 0.15
+        }
+      },
+      "ece_after": 0.06769,
+      "ece_before": 0.071107,
       "n": 500,
       "temperature": 0.75
     },
     "score": {
-      "ece_after": 0.177829,
-      "ece_before": 0.29653,
+      "by_language": {
+        "de": {
+          "ece_after": 0.128129,
+          "ece_before": 0.313627,
+          "n": 83,
+          "temperature": 0.15
+        },
+        "es": {
+          "ece_after": 0.160055,
+          "ece_before": 0.25,
+          "n": 84,
+          "temperature": 0.25
+        },
+        "fr": {
+          "ece_after": 0.103257,
+          "ece_before": 0.178855,
+          "n": 83,
+          "temperature": 0.25
+        },
+        "it": {
+          "ece_after": 0.118342,
+          "ece_before": 0.160282,
+          "n": 83,
+          "temperature": 0.75
+        },
+        "nl": {
+          "ece_after": 0.112712,
+          "ece_before": 0.359823,
+          "n": 83,
+          "temperature": 0.15
+        },
+        "pt": {
+          "ece_after": 0.102839,
+          "ece_before": 0.213012,
+          "n": 84,
+          "temperature": 0.25
+        }
+      },
+      "ece_after": 0.094548,
+      "ece_before": 0.233416,
       "n": 500,
       "temperature": 0.25
     }