Blame for host/clef_litert/head.py

"""The checkpoint's joint schema head (joint_schema_model.py `JointSchemaHead.forward`, one record) in NumPy float32.

Weights: joint_head.safetensors (bfloat16, 122 tensors, widened to float32 exactly; `load_state_dict(strict=True)` names)
and joint_head_config.json {hidden_size 4096, width 1024, routing_layers 2, layers 4, heads 16, feedforward 4096}.
Inputs: the row graph's hidden states at the record's T real positions (after the final RMSNorm), the record's ids
(processor form) and spans, and the lm_head table (`lm_head.weight` rows, for the lexical vectors).

    x        = hidden_norm(hidden)                          LayerNorm(4096); memory = memory_projection(x) [T, 1024]
    global   = x[T - 1];  q_i = mean of x over question i's span;  ctx_o = mean of x over option o's span
    lexical_o = mean of the lm_head rows of option o's span ids              (not normalized)
    queries_o = option_context_projection(ctx_o) + option_lexical_projection(lexical_o)
                + option_question_projection(q_i)                               i = the question of option o
    2 x evidence routing: queries += MHA(query_norm(queries), memory_norm(memory)); queries += FFN(feedforward_norm(.))
                          (pre-LN, 16 heads of 64, FFN = Linear 1024->4096, GELU, Linear 4096->1024)
    per question: w = softmax over its options of routed . question_projection(q_i) / 32;  summary_i = sum w * routed
    fields   = question_projection(q) + option_summary_norm(summary) + global_projection(global) + type_embedding(type)
    4 x TransformerDecoderLayer (norm_first, GELU): self-attention over the questions, cross-attention to the memory
              (not normalized), FFN;  fields = field_norm(fields)
    prior_o  = e^min(prior_logit_scale, ln 100) * cos(lexical_o, q_i + global)
    joint_o  = e^min(joint_logit_scale, ln 100) * cos(field_i, option_norm(routed_o)) + residual_scorer([f, o, f*o, |f-o|])
    logit_o  = prior_o + sigmoid(residual_gate) * joint_o

Every tensor is float32 as in the author's float32 module; LayerNorm eps 1e-5; attention = softmax(q k^T / 8) v per
head in float32; GELU = the exact erf form (evaluated in float64, rounded to float32); F.normalize eps 1e-12 and
F.cosine_similarity eps 1e-8 (each side divided by max(norm, eps), then the dot product, as torch computes it).
Dropout is the identity (eval). The per-question softmax of the logits is the response's (response.py)."""
from __future__ import annotations

import json
import math
from pathlib import Path

import numpy as np

from ._erf import gelu
from .tables import load_bf16_tensors

CONFIG = {"hidden_size": 4096, "width": 1024, "routing_layers": 2, "layers": 4, "heads": 16, "feedforward": 4096}
LN_EPS = 1e-5
F32 = np.float32


def layer_norm(x: np.ndarray, w: np.ndarray, b: np.ndarray, eps: float = LN_EPS) -> np.ndarray:
    mean = x.mean(axis=-1, keepdims=True, dtype=F32)
    xc = x - mean
    var = (xc * xc).mean(axis=-1, keepdims=True, dtype=F32)
    return (xc * (F32(1.0) / np.sqrt(var + F32(eps)))) * w + b


def normalize(x: np.ndarray, eps: float) -> np.ndarray:
    """x / max(||x||_2, eps) over the last axis (F.normalize)."""
    n = np.sqrt((x * x).sum(axis=-1, keepdims=True, dtype=F32))
    return x / np.maximum(n, F32(eps))


def softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
    e = np.exp(x - x.max(axis=axis, keepdims=True))
    return e / e.sum(axis=axis, keepdims=True, dtype=F32)


class _Linear:
    def __init__(self, w: np.ndarray, b: np.ndarray | None = None):
        self.wt = np.ascontiguousarray(w.T, dtype=F32)        # [in, out]
        self.b = None if b is None else np.ascontiguousarray(b, dtype=F32)

    def __call__(self, x: np.ndarray) -> np.ndarray:
        y = x @ self.wt
        return y if self.b is None else y + self.b


class _LayerNorm:
    def __init__(self, w: np.ndarray, b: np.ndarray):
        self.w, self.b = w.astype(F32), b.astype(F32)

    def __call__(self, x: np.ndarray) -> np.ndarray:
        return layer_norm(x, self.w, self.b)


class _Attention:
    """nn.MultiheadAttention (batch_first, eval, no masks) for one sequence: q_in [Lq, E], kv_in [Lk, E] -> [Lq, E]."""

    def __init__(self, t: dict, prefix: str, heads: int):
        w, b = t[prefix + "in_proj_weight"], t[prefix + "in_proj_bias"]
        e = w.shape[1]
        self.q = _Linear(w[:e], b[:e])
        self.k = _Linear(w[e:2 * e], b[e:2 * e])
        self.v = _Linear(w[2 * e:], b[2 * e:])
        self.out = _Linear(t[prefix + "out_proj.weight"], t[prefix + "out_proj.bias"])
        self.heads, self.dim = heads, e // heads
        self.scale = F32(1.0 / math.sqrt(self.dim))

    def __call__(self, q_in: np.ndarray, kv_in: np.ndarray) -> np.ndarray:
        h, d = self.heads, self.dim
        q = self.q(q_in).reshape(q_in.shape[0], h, d).transpose(1, 0, 2)          # [H, Lq, D]
        k = self.k(kv_in).reshape(kv_in.shape[0], h, d).transpose(1, 0, 2)        # [H, Lk, D]
        v = self.v(kv_in).reshape(kv_in.shape[0], h, d).transpose(1, 0, 2)
        p = softmax((q @ k.transpose(0, 2, 1)) * self.scale)                     # [H, Lq, Lk]
        o = (p @ v).transpose(1, 0, 2).reshape(q_in.shape[0], h * d)
        return self.out(o)


class _Tracked(dict):
    """The weights by name, recording which names were read (load_state_dict(strict=True): none missing, none unused)."""

    def __init__(self, *a):
        super().__init__(*a)
        self.read = set()

    def __getitem__(self, key):
        self.read.add(key)
        return super().__getitem__(key)


class JointHead:
    """The author's JointSchemaHead for one record. `scores(hidden, ids, questions, lm_head)` -> one float32 logit
    array per question, in the author's option order."""

    def __init__(self, tensors: dict[str, np.ndarray], config: dict | None = None):
        cfg = dict(CONFIG if config is None else config)
        if cfg != CONFIG:
            raise ValueError(f"head config {cfg} is not the release's {CONFIG}")
        self.config = cfg
        t = _Tracked({k: np.asarray(v, dtype=F32) for k, v in tensors.items()})
        heads = cfg["heads"]
        ln = lambda p: _LayerNorm(t[p + ".weight"], t[p + ".bias"])                # noqa: E731
        lin = lambda p, bias=True: _Linear(t[p + ".weight"], t[p + ".bias"] if bias else None)  # noqa: E731
        self.hidden_norm = ln("hidden_norm")
        self.memory_projection = lin("memory_projection", False)
        self.question_projection = lin("question_projection", False)
        self.option_question_projection = lin("option_question_projection", False)
        self.global_projection = lin("global_projection", False)
        self.option_context_projection = lin("option_context_projection", False)
        self.option_lexical_projection = lin("option_lexical_projection", False)
        self.type_embedding = t["type_embedding.weight"]
        self.evidence = [{"query_norm": ln(f"evidence_layers.{i}.query_norm"),
                          "memory_norm": ln(f"evidence_layers.{i}.memory_norm"),
                          "attention": _Attention(t, f"evidence_layers.{i}.attention.", heads),
                          "feedforward_norm": ln(f"evidence_layers.{i}.feedforward_norm"),
                          "ff0": lin(f"evidence_layers.{i}.feedforward.0"),
                          "ff3": lin(f"evidence_layers.{i}.feedforward.3")}
                         for i in range(cfg["routing_layers"])]
        self.option_summary_norm = ln("option_summary_norm")
        self.layers = [{"norm1": ln(f"layers.{i}.norm1"), "norm2": ln(f"layers.{i}.norm2"),
                        "norm3": ln(f"layers.{i}.norm3"),
                        "self_attn": _Attention(t, f"layers.{i}.self_attn.", heads),
                        "multihead_attn": _Attention(t, f"layers.{i}.multihead_attn.", heads),
                        "linear1": lin(f"layers.{i}.linear1"), "linear2": lin(f"layers.{i}.linear2")}
                       for i in range(cfg["layers"])]
        self.field_norm = ln("field_norm")
        self.option_norm = ln("option_norm")
        self.scorer0 = lin("residual_scorer.0")
        self.scorer3 = lin("residual_scorer.3")
        cap = F32(math.log(100.0))
        self.prior_scale = np.exp(np.minimum(t["prior_logit_scale"].reshape(()), cap)).astype(F32)
        self.joint_scale = np.exp(np.minimum(t["joint_logit_scale"].reshape(()), cap)).astype(F32)
        g = t["residual_gate"].reshape(())
        self.gate = (F32(1.0) / (F32(1.0) + np.exp(-g))).astype(F32)
        unused = set(t) - t.read                     # a missing name already raised KeyError when it was read
        if unused:
            raise ValueError(f"head weights not used: {sorted(unused)}")
        self.n_tensors = len(t)
        self.n_params = int(sum(v.size for v in dict.values(t)))

    @classmethod
    def from_files(cls, weights: str | Path, config: str | Path | None = None) -> "JointHead":
        cfg = json.loads(Path(config).read_text()) if config is not None else None
        return cls(load_bf16_tensors(weights), cfg)

    @staticmethod
    def lexical_rows(lm_head, ids, questions) -> list[np.ndarray]:
        """Per question, [n_options, 4096]: the mean of the lm_head rows of each option span's ids (float32)."""
        ids = np.asarray(ids, dtype=np.int64)
        out = []
        for q in questions:
            rows = [lm_head.rows(ids[a:b]).mean(axis=0, dtype=F32) for a, b in q.option_spans]
            out.append(np.stack(rows).astype(F32))
        return out

    def scores(self, hidden: np.ndarray, ids, questions, lm_head=None, lexical: list | None = None) -> list[np.ndarray]:
        """hidden [T, 4096] (the record's real positions), ids (len T), questions (request.EncodedQuestion or
        anything with question_span / option_spans / type_id) -> [float32 logits per question].
        `lexical` (one [n_options, 4096] array per question) replaces the lm_head gather (controls)."""
        hidden = np.asarray(hidden, dtype=F32)
        if hidden.ndim != 2 or hidden.shape[0] != len(ids):
            raise ValueError(f"hidden {hidden.shape} for {len(ids)} ids")
        if lexical is None:
            lexical = self.lexical_rows(lm_head, ids, questions)
        x = self.hidden_norm(hidden)                                            # [T, 4096]
        memory = self.memory_projection(x)                                      # [T, 1024]
        global_vector = x[-1]
        question_vectors = np.stack([x[a:b].mean(axis=0, dtype=F32) for a, b in (q.question_span for q in questions)])
        type_rows = self.type_embedding[[q.type_id for q in questions]]
        counts = [len(q.option_spans) for q in questions]
        context = [np.stack([x[a:b].mean(axis=0, dtype=F32) for a, b in q.option_spans]) for q in questions]

        queries = np.concatenate([
            self.option_context_projection(ctx) + self.option_lexical_projection(lex)
            + self.option_question_projection(question_vectors[i])[None]
            for i, (ctx, lex) in enumerate(zip(context, lexical))])             # [O, 1024]
        for layer in self.evidence:
            routed = layer["attention"](layer["query_norm"](queries), layer["memory_norm"](memory))
            queries = queries + routed
            queries = queries + layer["ff3"](gelu(layer["ff0"](layer["feedforward_norm"](queries))))
        split = np.split(queries, np.cumsum(counts)[:-1])

        base = self.question_projection(question_vectors)                       # [Q, 1024]
        summaries = []
        for field, options in zip(base, split):
            w = softmax((options @ field) / F32(math.sqrt(options.shape[-1])), axis=0)
            summaries.append((w[:, None] * options).sum(axis=0, dtype=F32))
        fields = (base + self.option_summary_norm(np.stack(summaries)) + self.global_projection(global_vector)[None]
                  + type_rows)
        for layer in self.layers:
            f1 = layer["norm1"](fields)
            fields = fields + layer["self_attn"](f1, f1)
            fields = fields + layer["multihead_attn"](layer["norm2"](fields), memory)
            fields = fields + layer["linear2"](gelu(layer["linear1"](layer["norm3"](fields))))
        fields = self.field_norm(fields)

        out = []
        for i, (field, lex, routed) in enumerate(zip(fields, lexical, split)):
            anchor = normalize(question_vectors[i] + global_vector, 1e-12)
            prior = self.prior_scale * (normalize(lex, 1e-12) @ anchor)
            options = self.option_norm(routed)
            f = np.broadcast_to(field, options.shape)
            cosine = (normalize(f, 1e-8) * normalize(options, 1e-8)).sum(axis=-1, dtype=F32)
            features = np.concatenate([f, options, f * options, np.abs(f - options)], axis=-1)
            residual = self.scorer3(gelu(self.scorer0(features)))[:, 0]
            joint = self.joint_scale * cosine + residual
            out.append((prior + self.gate * joint).astype(F32))
        return out