Re-export: bake softmax/sigmoid into the graph (4 named probability outputs); external-data weights + export manifest

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+alignscore-large.onnx.data filter=lfs diff=lfs merge=lfs -text

README.md

@@ -34,15 +34,18 @@ general-purpose defaults.
 
 ## What this is
 
-This is a faithful ONNX export of the encoder + pooler + the two heads used for
-alignment scoring:
+A faithful ONNX export of the encoder + pooler + the two alignment-scoring heads,
+with the output activations **baked into the graph** so the model emits the four
+probabilities directly:
 
 - a RoBERTa-large encoder with pooling layer (`pooler_output = tanh(dense(h[:,0]))`)
-- `tri_layer`: `Linear(hidden, 3)` -> raw 3-way logits
-- `reg_layer`: `Linear(hidden, 1)` -> raw regression logit
+- `tri_layer`: `Linear(hidden, 3)` -> `softmax` -> 3-way probabilities
+- `reg_layer`: `Linear(hidden, 1)` -> `sigmoid` -> alignment probability
 
-Activations are **not** baked into the graph: it emits raw logits, and the
-caller applies `softmax` to the 3-way head and `sigmoid` to the regression head.
+This is the key difference from the earlier revision of this repo, which emitted
+raw `tri_logits` / `reg_logit` and left `softmax`/`sigmoid` to the caller. Baking
+the activations in makes the graph self-contained: a single direction scores a
+`(context, claim)` pair straight to probabilities, no post-processing.
 
 ### Graph I/O
 
@@ -50,28 +53,36 @@ caller applies `softmax` to the 3-way head and `sigmoid` to the regression head.
 |--------|-----------|------|-------|
 | `input_ids` | input | int64 | `[batch, seq]` (dynamic) |
 | `attention_mask` | input | int64 | `[batch, seq]` (dynamic) |
-| `tri_logits` | output | float32 | `[batch, 3]` -> softmax -> `[p_aligned, p_neutral, p_contradict]` |
-| `reg_logit` | output | float32 | `[batch, 1]` -> sigmoid -> `p_aligned_reg` |
+| `p_aligned_3way` | output | float32 | `[batch]` (softmax over tri head, aligned) |
+| `p_neutral_3way` | output | float32 | `[batch]` (softmax over tri head, neutral) |
+| `p_contradict_3way` | output | float32 | `[batch]` (softmax over tri head, contradict) |
+| `p_aligned_reg` | output | float32 | `[batch]` (sigmoid over reg head) |
 
 There is no `token_type_ids` input: a sentence pair is encoded into a single
 `input_ids` sequence with `</s></s>` separators, exactly as
 `AutoTokenizer("roberta-large")(context, claim)` produces. The bundled
 `tokenizer.json` is the matching fast tokenizer. Use `max_length=512`.
 
-Opset 17. Exported with PyTorch (TorchScript exporter).
+Opset 17. Weights are stored as ONNX external data in `alignscore-large.onnx.data`
+(the `.onnx` is the graph; keep the two files side by side). The
+`alignscore-export-manifest.json` records the source checkpoint SHA-256, the
+upstream HF revision, the opset, and the exact input/output tensor names.
 
 ## Files
 
-- `alignscore-large.onnx` - the model (~1.36 GB)
+- `alignscore-large.onnx` - the model graph (~0.2 MB)
+- `alignscore-large.onnx.data` - external weights (~1.4 GB); must sit next to the `.onnx`
+- `alignscore-export-manifest.json` - checkpoint SHA-256, HF revision, opset, I/O names
 - `tokenizer.json` - roberta-large fast tokenizer with the pair post-processor
 
 ## Parity
 
 Verified against the original PyTorch model's scores on a 136-pair corpus:
-**max absolute difference 1.18e-07** across both directions and all four
-probabilities. The source checkpoint SHA-256 is asserted equal to the reference
-before export, so these are provably the same weights. A high-confidence
-contradiction pair scores `p_contradict_3way` 0.982 / 0.977 (forward / reverse).
+**max absolute difference 5e-06** on `p_contradict_3way` and **0** on
+`p_aligned_reg`, across both directions, with **zero** verdict flips through a
+downstream 0.75 bidirectional contradiction gate. The source checkpoint SHA-256
+is asserted equal to the reference before export, so these are provably the same
+weights.
 
 ## Usage (ONNX Runtime, Python)
 
@@ -80,20 +91,21 @@ import numpy as np, onnxruntime as ort
 from tokenizers import Tokenizer
 
 tok = Tokenizer.from_file("tokenizer.json")
-sess = ort.InferenceSession("alignscore-large.onnx")
+sess = ort.InferenceSession("alignscore-large.onnx")  # loads .onnx.data automatically
 
 def score(context, claim):
     enc = tok.encode(context, claim)
     ids = np.array([enc.ids], dtype=np.int64)
     mask = np.array([enc.attention_mask], dtype=np.int64)
-    tri, reg = sess.run(["tri_logits", "reg_logit"],
-                        {"input_ids": ids, "attention_mask": mask})
-    e = np.exp(tri[0] - tri[0].max()); tri_p = e / e.sum()
+    pa, pn, pc, pr = sess.run(
+        ["p_aligned_3way", "p_neutral_3way", "p_contradict_3way", "p_aligned_reg"],
+        {"input_ids": ids, "attention_mask": mask},
+    )
     return {
-        "p_aligned_3way": float(tri_p[0]),
-        "p_neutral_3way": float(tri_p[1]),
-        "p_contradict_3way": float(tri_p[2]),
-        "p_aligned_reg": float(1 / (1 + np.exp(-reg[0][0]))),
+        "p_aligned_3way": float(pa[0]),
+        "p_neutral_3way": float(pn[0]),
+        "p_contradict_3way": float(pc[0]),
+        "p_aligned_reg": float(pr[0]),
     }
 ```
 

alignscore-export-manifest.json

@@ -0,0 +1,18 @@
+{
+  "schema_version": 1,
+  "checkpoint_sha256": "ff4336312b377edcbcdad5694a2d09d73dc4225422c0422d810aa7e78485e32d",
+  "hf_revision": "8509e78d25bb914939fc585c626500c9b2944249",
+  "opset": 17,
+  "encoder_name": "roberta-large",
+  "input_names": [
+    "input_ids",
+    "attention_mask"
+  ],
+  "output_names": [
+    "p_aligned_3way",
+    "p_neutral_3way",
+    "p_contradict_3way",
+    "p_aligned_reg"
+  ],
+  "generated_at": "2026-07-05T03:01:18Z"
+}

alignscore-large.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:350be954d95762c12157682aae77cd84699e189f6752d0caf8d38298b2ef2f90
-size 1421893132
+oid sha256:7c4a95d87043d6a543d5089d3363ac0c3168bbfee93be06dadb7c36e9d9d3c76
+size 210708

alignscore-large.onnx.data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c842f392dae6a5c777bb9792332731a14d53719815afac394f1ea1f0ce29c03
+size 1421578240

tokenizer.json

@@ -50,7 +50,7 @@
       "single_word": false,
       "lstrip": true,
       "rstrip": false,
-      "normalized": true,
+      "normalized": false,
       "special": true
     }
   ],