js/src/model.ts

/**
 * InfonHeadsModel — the main public class.
 *
 * Pipeline:
 *   text → tokenize → backbone CLS → multi-head classifier → 6 labels
 *
 * Same architectural shape as ``@cp500/infon-coref`` (one big ONNX
 * for the backbone, one tiny ONNX for the heads), so consumers using
 * both packages can share the runtime adapter, hub fetcher, and
 * tokenizer infra at the npm-resolution level. We don't import from
 * sibling packages to avoid coupling — the duplication is small and
 * keeps each package independently versioned.
 */

import { fetchHubFile, fetchHubJson } from './hub.js';
import {
  CONDITIONAL_LABELS,
  DIRECTION_LABELS,
  ONNX_OUTPUT_NAMES,
  POLARITY_LABELS,
  RELATION_LABELS,
  SPATIAL_LABELS,
  TENSE_LABELS,
} from './labels.js';
import {
  createSession,
  makeTensor,
  type OrtSession,
  type OrtTensor,
} from './ort.js';
import { argmaxWithProb } from './softmax.js';
import { loadTokenizer } from './tokenizer.js';
import type { ClassifyResult, LoadedModel, ModelOptions } from './types.js';

/** Files referenced by the HF model repo's ``meta.json``. */
interface RepoMeta {
  hidden_size: number;
  max_length: number;
  cls_token: string;
  files: {
    backbone_fp32?: string;
    backbone_fp16?: string;
    classifier_fp32?: string;
    classifier_fp16?: string;
    tokenizer: string;
  };
}

/**
 * Multilingual single-sentence linguistic classifiers.
 *
 * @example Browser
 * ```ts
 * import { InfonHeadsModel } from '@cp500/infon-heads';
 *
 * const model = await InfonHeadsModel.fromHub('cp500/infon-heads', {
 *   precision: 'fp16',
 * });
 *
 * const r = await model.classify(
 *   'Toyota did not raise battery output last quarter because demand fell.'
 * );
 * console.log(r.polarity);             // 'negated'
 * console.log(r.tense);                // 'past'
 * console.log(r.relationType);         // 'causal'
 * console.log(r.comparativeDirection); // 'decrease'
 * console.log(r.confidence.polarity);  // 0.94
 * ```
 *
 * @example Batched
 * ```ts
 * const results = await model.classifyMany([
 *   'BAE relocated its drone assembly to Sydney.',
 *   '丰田预计明年销量将增长12%。',
 * ]);
 * ```
 */
export class InfonHeadsModel {
  private constructor(private readonly loaded: LoadedModel) {}

  /** Load model artefacts from a Hugging Face repo. Caches downloads
   * in the browser Cache API; subsequent loads reuse the disk copy. */
  static async fromHub(
    repo: string,
    opts: ModelOptions & { revision?: string } = {},
  ): Promise<InfonHeadsModel> {
    const meta = await fetchHubJson<RepoMeta>({
      repo,
      path: 'meta.json',
      revision: opts.revision,
    });
    const precision = opts.precision ?? 'fp16';
    const backboneFile =
      precision === 'fp16'
        ? meta.files.backbone_fp16 ?? meta.files.backbone_fp32
        : meta.files.backbone_fp32 ?? meta.files.backbone_fp16;
    const classifierFile =
      precision === 'fp16'
        ? meta.files.classifier_fp16 ?? meta.files.classifier_fp32
        : meta.files.classifier_fp32 ?? meta.files.classifier_fp16;
    if (!backboneFile || !classifierFile) {
      throw new Error(
        `repo ${repo} is missing backbone/classifier ONNX files; check meta.json`,
      );
    }

    const [backboneBuf, classifierBuf, tokBuf] = await Promise.all([
      fetchHubFile({ repo, path: backboneFile, revision: opts.revision }),
      fetchHubFile({
        repo,
        path: classifierFile,
        revision: opts.revision,
      }),
      fetchHubFile({
        repo,
        path: meta.files.tokenizer,
        revision: opts.revision,
      }),
    ]);
    if (opts.debug) {
      const mb = (b: ArrayBuffer) => (b.byteLength / 1e6).toFixed(1);
      console.debug(
        `[infon-heads] loaded backbone ${mb(backboneBuf)} MB, ` +
          `classifier ${mb(classifierBuf)} MB, ` +
          `tokenizer ${mb(tokBuf)} MB`,
      );
    }
    return InfonHeadsModel.#fromBuffers(
      backboneBuf,
      classifierBuf,
      tokBuf,
      meta,
      opts,
    );
  }

  /** Load from a local directory (Node fs path or a browser URL). */
  static async fromLocal(
    baseUrl: string,
    opts: ModelOptions = {},
  ): Promise<InfonHeadsModel> {
    const isPath =
      typeof window === 'undefined' && !baseUrl.startsWith('http');
    const join = (name: string) =>
      isPath
        ? `${baseUrl.replace(/\/$/, '')}/${name}`
        : new URL(name, baseUrl.endsWith('/') ? baseUrl : baseUrl + '/')
            .href;
    const meta = await loadJson<RepoMeta>(join('meta.json'), isPath);
    const precision = opts.precision ?? 'fp16';
    const backboneFile =
      precision === 'fp16'
        ? meta.files.backbone_fp16 ?? meta.files.backbone_fp32
        : meta.files.backbone_fp32 ?? meta.files.backbone_fp16;
    const classifierFile =
      precision === 'fp16'
        ? meta.files.classifier_fp16 ?? meta.files.classifier_fp32
        : meta.files.classifier_fp32 ?? meta.files.classifier_fp16;
    if (!backboneFile || !classifierFile) {
      throw new Error('local model missing backbone/classifier ONNX files');
    }
    const [bbBuf, clBuf, tkBuf] = await Promise.all([
      loadBytes(join(backboneFile), isPath),
      loadBytes(join(classifierFile), isPath),
      loadBytes(join(meta.files.tokenizer), isPath),
    ]);
    return InfonHeadsModel.#fromBuffers(bbBuf, clBuf, tkBuf, meta, opts);
  }

  static async #fromBuffers(
    backboneBuf: ArrayBuffer,
    classifierBuf: ArrayBuffer,
    tokBuf: ArrayBuffer,
    meta: RepoMeta,
    opts: ModelOptions,
  ): Promise<InfonHeadsModel> {
    const device = opts.device ?? 'auto';
    const [backbone, classifier, tokenizer] = await Promise.all([
      createSession(new Uint8Array(backboneBuf), device),
      createSession(new Uint8Array(classifierBuf), device),
      loadTokenizer(new Uint8Array(tokBuf)),
    ]);
    return new InfonHeadsModel({
      backbone,
      classifier,
      tokenizer,
      meta: {
        hiddenSize: meta.hidden_size,
        maxLength: opts.maxLength ?? meta.max_length,
        precision: opts.precision ?? 'fp16',
        device:
          device === 'auto'
            ? typeof window !== 'undefined'
              ? 'auto-browser'
              : 'auto-node'
            : device,
      },
    });
  }

  /** Classify a single sentence. */
  async classify(
    text: string,
    opts: ModelOptions = {},
  ): Promise<ClassifyResult> {
    const results = await this.classifyMany([text], opts);
    return results[0];
  }

  /** Classify a batch of sentences. Single backbone forward over the
   * padded batch; one ``classifier`` call per sentence (the heads
   * are tiny so per-row execution dominates only the network
   * round-trip in the browser — negligible).
   *
   * If you have many short sentences, batching here is much faster
   * than calling ``classify`` in a loop because the backbone forward
   * is parallel-friendly while the per-sentence loop forces a
   * sequential ORT round trip per item.
   */
  async classifyMany(
    texts: string[],
    opts: ModelOptions = {},
  ): Promise<ClassifyResult[]> {
    const debug = opts.debug ?? false;
    const t0 = nowMs();

    if (texts.length === 0) return [];

    // 1. Tokenize (per-sentence, then pad to common length).
    const max = opts.maxLength ?? this.loaded.meta.maxLength;
    const encs = texts.map((t) =>
      this.loaded.tokenizer.tokenize(t, { maxLength: max }),
    );
    const T = Math.max(...encs.map((e) => e.inputIds.length));

    const ids = new BigInt64Array(texts.length * T);
    const mask = new BigInt64Array(texts.length * T);
    for (let i = 0; i < texts.length; i++) {
      const e = encs[i];
      for (let t = 0; t < e.inputIds.length; t++) {
        ids[i * T + t] = e.inputIds[t];
        mask[i * T + t] = e.attentionMask[t];
      }
      // Tail of row i is already zeroed (PAD), which is correct since
      // attention_mask=0 there.
    }
    const t1 = nowMs();

    // 2. Backbone forward → CLS (B, H).
    const idsT = await makeTensor('int64', ids, [texts.length, T]);
    const maskT = await makeTensor('int64', mask, [texts.length, T]);
    const bbOut = await this.loaded.backbone.run({
      input_ids: idsT,
      attention_mask: maskT,
    });
    const clsTensor = bbOut.cls ?? bbOut.last_hidden_state;
    if (!clsTensor) {
      throw new Error(
        `backbone outputs missing 'cls'; got [${Object.keys(bbOut).join(', ')}]`,
      );
    }
    const t2 = nowMs();

    // 3. Classifier forward — accepts (B, H), emits 6 logit tensors
    // each of shape (B, n_classes_for_that_head).
    const clsFlat = floatArray(clsTensor);
    const H = this.loaded.meta.hiddenSize;
    if (clsFlat.length !== texts.length * H) {
      throw new Error(
        `cls shape mismatch: expected ${texts.length}*${H}=${
          texts.length * H
        }, got ${clsFlat.length}`,
      );
    }
    const clsT = await makeTensor('float32', clsFlat, [texts.length, H]);
    const headOut = await this.loaded.classifier.run({ cls: clsT });
    const t3 = nowMs();

    // 4. Argmax + softmax per sentence per head.
    const logits = ONNX_OUTPUT_NAMES.map((name) => {
      const t = headOut[name];
      if (!t) {
        throw new Error(
          `classifier output '${name}' missing; got [${Object.keys(
            headOut,
          ).join(', ')}]`,
        );
      }
      return floatArray(t);
    });

    const out: ClassifyResult[] = [];
    for (let i = 0; i < texts.length; i++) {
      const polarityRow = sliceRow(logits[0], i, POLARITY_LABELS.length);
      const tenseRow = sliceRow(logits[1], i, TENSE_LABELS.length);
      const conditionalRow = sliceRow(
        logits[2],
        i,
        CONDITIONAL_LABELS.length,
      );
      const relationRow = sliceRow(logits[3], i, RELATION_LABELS.length);
      const spatialRow = sliceRow(logits[4], i, SPATIAL_LABELS.length);
      const directionRow = sliceRow(logits[5], i, DIRECTION_LABELS.length);

      const pol = argmaxWithProb(polarityRow);
      const tense = argmaxWithProb(tenseRow);
      const cond = argmaxWithProb(conditionalRow);
      const rel = argmaxWithProb(relationRow);
      const spat = argmaxWithProb(spatialRow);
      const dir = argmaxWithProb(directionRow);

      out.push({
        text: texts[i],
        polarity: POLARITY_LABELS[pol.idx],
        tense: TENSE_LABELS[tense.idx],
        conditional: CONDITIONAL_LABELS[cond.idx] === 'yes',
        relationType: RELATION_LABELS[rel.idx],
        spatialRelation: SPATIAL_LABELS[spat.idx],
        comparativeDirection: DIRECTION_LABELS[dir.idx],
        confidence: {
          polarity: pol.prob,
          tense: tense.prob,
          conditional: cond.prob,
          relationType: rel.prob,
          spatialRelation: spat.prob,
          comparativeDirection: dir.prob,
        },
        timing: {
          tokenize: t1 - t0,
          backbone: t2 - t1,
          heads: t3 - t2,
          total: t3 - t0,
        },
      });
    }

    if (debug) {
      console.debug('[infon-heads] timings (ms)', {
        n: texts.length,
        tokenize: t1 - t0,
        backbone: t2 - t1,
        heads: t3 - t2,
        total: t3 - t0,
      });
    }

    return out;
  }

  /** Architecture metadata loaded from the repo's ``meta.json``. */
  get meta() {
    return this.loaded.meta;
  }
}

// ── Internal helpers ────────────────────────────────────────────────

function nowMs(): number {
  return typeof performance !== 'undefined'
    ? performance.now()
    : Date.now();
}

function sliceRow(
  flat: Float32Array,
  rowIdx: number,
  cols: number,
): Float32Array {
  return flat.subarray(rowIdx * cols, (rowIdx + 1) * cols);
}

function floatArray(t: OrtTensor): Float32Array {
  if (t.data instanceof Float32Array) return t.data;
  if (t.data instanceof Uint16Array) {
    // FP16 → FP32 (matches ``@cp500/infon-coref``).
    const out = new Float32Array(t.data.length);
    for (let i = 0; i < t.data.length; i++) {
      out[i] = halfToFloat(t.data[i]);
    }
    return out;
  }
  throw new Error(
    `expected Float32Array or Uint16Array tensor, got ${
      (t.data as { constructor: { name: string } }).constructor.name
    }`,
  );
}

function halfToFloat(h: number): number {
  const sign = (h & 0x8000) >> 15;
  const exp = (h & 0x7c00) >> 10;
  const frac = h & 0x03ff;
  if (exp === 0) {
    return (sign ? -1 : 1) * Math.pow(2, -14) * (frac / 1024);
  } else if (exp === 0x1f) {
    return frac ? NaN : (sign ? -1 : 1) * Infinity;
  }
  return (sign ? -1 : 1) * Math.pow(2, exp - 15) * (1 + frac / 1024);
}

async function loadJson<T>(path: string, isFsPath: boolean): Promise<T> {
  if (isFsPath) {
    const fs = await import('node:fs/promises');
    return JSON.parse(await fs.readFile(path, 'utf-8')) as T;
  }
  const r = await fetch(path);
  if (!r.ok) throw new Error(`fetch ${path}: ${r.status}`);
  return (await r.json()) as T;
}

async function loadBytes(
  path: string,
  isFsPath: boolean,
): Promise<ArrayBuffer> {
  if (isFsPath) {
    const fs = await import('node:fs/promises');
    const buf = await fs.readFile(path);
    return buf.buffer.slice(
      buf.byteOffset,
      buf.byteOffset + buf.byteLength,
    );
  }
  const r = await fetch(path);
  if (!r.ok) throw new Error(`fetch ${path}: ${r.status}`);
  return await r.arrayBuffer();
}