Upload js/src/pairs.ts with huggingface_hub

js/src/pairs.ts

@@ -0,0 +1,145 @@
+/**
+ * Pair-index builder + per-mention argmax + cluster grouping.
+ *
+ * Pure JS, no ORT, no DOM. Mirrors the Python helpers in
+ * ``infon/scripts/coref_onnx_experiment.py`` (``build_pairs`` /
+ * ``split_pairs_by_mention``) so the JS pipeline produces
+ * bit-identical pair tensors.
+ */
+
+/**
+ * Enumerate ``(i, j)`` candidate pairs for ``M`` mentions.
+ *
+ * For mention ``m`` (1-indexed because index 0 is DUMMY) we emit
+ * ``(m, 0), (m, 1), …, (m, m-1)`` — DUMMY first, then every earlier
+ * mention. This is the same triangular shape the Python
+ * ``build_pairs`` returns; the scorer ONNX expects this exact layout
+ * because the in-graph ``index_select`` over the prepended DUMMY
+ * relies on j=0 meaning "no antecedent."
+ *
+ * @param nMentions  number of mentions in the doc
+ * @returns ``[pairI, pairJ]`` BigInt64 typed arrays of equal length.
+ *          Lengths: ``M*(M+1)/2``.
+ */
+export function buildPairs(nMentions: number): [BigInt64Array, BigInt64Array] {
+  const pi: bigint[] = [];
+  const pj: bigint[] = [];
+  for (let m = 1; m <= nMentions; m++) {
+    pi.push(BigInt(m));
+    pj.push(0n);
+    for (let j = 1; j < m; j++) {
+      pi.push(BigInt(m));
+      pj.push(BigInt(j));
+    }
+  }
+  return [BigInt64Array.from(pi), BigInt64Array.from(pj)];
+}
+
+/**
+ * Group flat pair scores back into per-mention argmax decisions.
+ * Mirrors ``split_pairs_by_mention`` in the Python harness.
+ *
+ * @returns ``decisions[i]`` = the mention index (1-based) chosen as
+ *          mention i's antecedent, or ``0`` for DUMMY (no antecedent).
+ *          Translate to 0-based mention indices with ``decisions[i] - 1``.
+ */
+export function pickAntecedents(
+  nMentions: number,
+  pairI: BigInt64Array,
+  pairJ: BigInt64Array,
+  scores: Float32Array,
+): { antecedent: number; score: number }[] {
+  const out: { antecedent: number; score: number }[] = [];
+  for (let m = 1; m <= nMentions; m++) {
+    let bestIdx = -1;
+    let bestScore = -Infinity;
+    for (let k = 0; k < pairI.length; k++) {
+      if (Number(pairI[k]) !== m) continue;
+      const s = scores[k];
+      if (s > bestScore) {
+        bestScore = s;
+        bestIdx = k;
+      }
+    }
+    out.push({
+      antecedent: bestIdx >= 0 ? Number(pairJ[bestIdx]) : 0,
+      score: bestScore,
+    });
+  }
+  return out;
+}
+
+/**
+ * Group antecedent decisions into clusters using union-find.
+ *
+ * Each mention either points to DUMMY (starts its own cluster) or to
+ * an earlier mention (joins that mention's cluster). Cluster IDs are
+ * dense 0-based; singletons are not assigned a cluster (returned as
+ * ``-1``) so callers can render them differently.
+ *
+ * @param decisions  ``decisions[i].antecedent`` is the *1-based*
+ *                   mention index this mention links to, or ``0`` for
+ *                   DUMMY. (Same convention as the model output.)
+ * @returns
+ *   - ``cluster[i]`` — cluster id for mention i, or -1 if singleton
+ *   - ``clusters`` — list of multi-mention clusters, each a list of
+ *     mention indices in document order
+ */
+export function groupClusters(
+  decisions: { antecedent: number }[],
+): { cluster: number[]; clusters: number[][] } {
+  const n = decisions.length;
+  // Union-find. parent[i] points to a smaller-or-equal mention index.
+  const parent = Array.from({ length: n }, (_, i) => i);
+  const find = (x: number): number => {
+    while (parent[x] !== x) {
+      parent[x] = parent[parent[x]];   // path compression
+      x = parent[x];
+    }
+    return x;
+  };
+  const union = (a: number, b: number) => {
+    const ra = find(a);
+    const rb = find(b);
+    if (ra !== rb) {
+      // Always attach the higher-index root under the lower-index root
+      // so cluster representatives are first-mention.
+      if (ra < rb) parent[rb] = ra;
+      else parent[ra] = rb;
+    }
+  };
+
+  for (let i = 0; i < n; i++) {
+    const ant = decisions[i].antecedent;
+    if (ant > 0) {
+      // ant is 1-based; the mention it points to is ant - 1.
+      union(i, ant - 1);
+    }
+  }
+
+  // Bucket by root.
+  const roots: number[][] = [];
+  const rootIdx = new Map<number, number>();
+  for (let i = 0; i < n; i++) {
+    const r = find(i);
+    let idx = rootIdx.get(r);
+    if (idx === undefined) {
+      idx = roots.length;
+      roots.push([]);
+      rootIdx.set(r, idx);
+    }
+    roots[idx].push(i);
+  }
+
+  // Collapse: only multi-mention clusters get a stable id; singletons
+  // get -1.
+  const cluster = new Array<number>(n).fill(-1);
+  const clusters: number[][] = [];
+  for (const group of roots) {
+    if (group.length < 2) continue;
+    const cid = clusters.length;
+    clusters.push(group);
+    for (const m of group) cluster[m] = cid;
+  }
+  return { cluster, clusters };
+}