FamiliarTools/AlignScore-large-onnx

AlignScore-large (ONNX)

An ONNX export of AlignScore-large (RoBERTa-large with AlignScore's 3-way and regression heads), for in-process inference without a Python/PyTorch runtime.

Why this exists

Upstream AlignScore ships only PyTorch Lightning checkpoints built from a custom BERTAlignModel module - no config.json, no safetensors - so optimum-cli export onnx cannot consume it, and no ONNX build existed. This is that build, so anyone who wants to experiment with AlignScore-large can, without standing up a PyTorch runtime or writing a custom export pathway.

It reflects a Familiar Tools belief: a specialized, right-sized model that runs efficiently and in-process beats reaching for a large, general, resource-hungry one. Exporting a focused model to ONNX is part of that - it makes the model cheap to run, easy to embed, and light on dependencies. Custom, deliberately engineered solutions tend to be more efficient and more resource-aware than general-purpose defaults.

What this is

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) -> softmax -> 3-way probabilities
• reg_layer: Linear(hidden, 1) -> sigmoid -> alignment probability

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

Tensor	Direction	Type	Shape
input_ids	input	int64	[batch, seq] (dynamic)
attention_mask	input	int64	[batch, seq] (dynamic)
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. 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 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 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)

import numpy as np, onnxruntime as ort
from tokenizers import Tokenizer

tok = Tokenizer.from_file("tokenizer.json")
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)
    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(pa[0]),
        "p_neutral_3way": float(pn[0]),
        "p_contradict_3way": float(pc[0]),
        "p_aligned_reg": float(pr[0]),
    }

License and attribution

Released under the MIT License, matching upstream.

• AlignScore: Zha et al., AlignScore: Evaluating Factual Consistency with a Unified Alignment Function, ACL 2023 (arXiv:2305.16739, code).
• Original weights: yzha/AlignScore (revision 8509e78d25bb914939fc585c626500c9b2944249).
• Base encoder: RoBERTa-large (Liu et al., 2019).

This repo redistributes a derivative (ONNX export) of the above under the same MIT terms. No weights were retrained or modified; only the inference graph was re-expressed.