trace-validation.js

/* alphaBell Model Observatory trace import checks. No data is changed here. */
(function (root) {
  'use strict';
  const LIMITS = Object.freeze({ frames: 256, context: 2048, tokenText: 8192, prompt: 4000 });
  const EPSILON = 0.00001;
  const object = value => value !== null && typeof value === 'object' && !Array.isArray(value);
  const number = value => typeof value === 'number' && Number.isFinite(value);
  const nonnegative = value => number(value) && value >= 0;
  const probability = value => nonnegative(value) && value <= 1 + EPSILON;
  const integer = (value, minimum, maximum) => Number.isInteger(value) && value >= minimum && value <= maximum;
  const text = value => typeof value === 'string' && value.length <= LIMITS.tokenText;
  function requireValue(condition, message) {
    if (!condition) throw new Error(message);
  }
  function vector(value, length, predicate, message) {
    requireValue(Array.isArray(value) && value.length === length && value.every(predicate), message);
  }
  function validateObservatoryTrace(d, models) {
    requireValue(object(d) && Array.isArray(models), 'This file is not a supported model trace.');
    requireValue(['flint.trace.v1', 'model.trace.v1'].includes(d.schema), 'This file is not a supported model trace.');
    const id = d.model?.id || (d.schema === 'flint.trace.v1' ? 'flint' : null);
    const model = models.find(item => item.id === id);
    requireValue(Boolean(model) && (d.schema !== 'flint.trace.v1' || id === 'flint'), 'The trace identifies an unsupported model.');
    const c = model.architecture;
    requireValue(object(c) && object(d.architecture), 'The trace has no supported architecture.');
    const requiredKeys = [
      'num_hidden_layers', 'num_attention_heads', 'num_key_value_heads', 'hidden_size',
      'head_dim', 'intermediate_size', 'layer_types'
    ];
    const defaultZero = ['num_experts', 'num_experts_per_tok', 'first_k_dense_replace'];
    const defaultArray = ['mlp_only_layers', 'no_rope_layer_indices'];
    const conditionalKeys = [
      'sliding_window', 'moe_intermediate_size', 'shared_expert_intermediate_size',
      'linear_num_key_heads', 'linear_num_value_heads', 'linear_key_head_dim', 'linear_value_head_dim'
    ];
    const mismatch = 'The trace architecture does not match ' + model.short_name + '.';
    for (const key of requiredKeys) {
      requireValue(JSON.stringify(d.architecture[key]) === JSON.stringify(c[key]), mismatch);
    }
    for (const key of defaultZero) requireValue((d.architecture[key] ?? 0) === (c[key] ?? 0), mismatch);
    for (const key of defaultArray) {
      requireValue(JSON.stringify(d.architecture[key] ?? []) === JSON.stringify(c[key] ?? []), mismatch);
    }
    for (const key of conditionalKeys) {
      if (c[key] !== undefined) requireValue(d.architecture[key] === c[key], mismatch);
    }
    if (d.model?.quantization !== undefined) {
      requireValue(d.model.quantization === model.quantization, 'The trace quantization does not match the selected model.');
    }
    requireValue(Array.isArray(d.frames) && d.frames.length > 0 && d.frames.length <= LIMITS.frames,
      'The trace must contain between 1 and 256 measured frames.');
    requireValue(object(d.start) && Array.isArray(d.start.tokens) && d.start.tokens.length > 0 &&
      d.start.tokens.length <= LIMITS.context && d.start.tokens.every(text), 'The trace has invalid prompt tokens.');
    if (d.start.prompt_tokens !== undefined) {
      requireValue(d.start.prompt_tokens === d.start.tokens.length, 'The trace prompt-token count is inconsistent.');
    }
    if (d.prompt !== undefined) {
      requireValue(typeof d.prompt === 'string' && d.prompt.length <= LIMITS.prompt, 'The trace prompt is invalid.');
    }
    if (d.settings !== undefined) {
      requireValue(object(d.settings), 'The trace settings are invalid.');
      const settings = d.settings;
      if (settings.max_tokens !== undefined) requireValue(integer(settings.max_tokens, 1, LIMITS.frames), 'The output limit is invalid.');
      if (settings.temperature !== undefined) requireValue(nonnegative(settings.temperature) && settings.temperature <= 1.5, 'The temperature is invalid.');
      if (settings.seed !== undefined) requireValue(integer(settings.seed, 0, 2147483647), 'The seed is invalid.');
      if (settings.reasoning !== undefined) requireValue(typeof settings.reasoning === 'boolean', 'The reasoning setting is invalid.');
    }
    let previousElapsed = -1;
    for (let frameIndex = 0; frameIndex < d.frames.length; frameIndex++) {
      const f = d.frames[frameIndex];
      requireValue(object(f) && text(f.text), 'The trace contains invalid token text.');
      requireValue(integer(f.input_position, 0, LIMITS.context - 1) &&
        f.input_position === d.start.tokens.length - 1 + frameIndex, 'The trace input position is inconsistent.');
      if (f.model !== undefined) requireValue(f.model === id, 'A trace frame identifies a different model.');
      if (f.step !== undefined) requireValue(f.step === frameIndex, 'The trace token order is inconsistent.');
      requireValue(integer(f.token_id, 0, (c.vocab_size || Number.MAX_SAFE_INTEGER) - 1), 'The trace contains an invalid token ID.');
      requireValue(nonnegative(f.duration) && nonnegative(f.elapsed) && f.elapsed >= previousElapsed,
        'The trace contains invalid token timings.');
      previousElapsed = f.elapsed;
      if (f.input_text !== undefined) requireValue(text(f.input_text), 'The trace contains invalid input-token text.');
      if (f.token_bytes !== undefined) {
        requireValue(Array.isArray(f.token_bytes) && f.token_bytes.length <= LIMITS.tokenText &&
          f.token_bytes.every(value => integer(value, 0, 255)), 'The trace contains invalid token bytes.');
      }
      requireValue(Array.isArray(f.layers) && f.layers.length === c.num_hidden_layers, 'The trace contains an incomplete frame.');
      for (let i = 0; i < c.num_hidden_layers; i++) {
        const l = f.layers[i];
        const linear = c.layer_types[i] === 'linear_attention';
        const heads = linear ? c.linear_num_value_heads : c.num_attention_heads;
        const moe = Boolean(c.num_experts) && !(c.mlp_only_layers || []).includes(i) && i >= (c.first_k_dense_replace || 0);
        requireValue(object(l), 'The trace contains an invalid layer.');
        if (l.index !== undefined) requireValue(l.index === i, 'The trace layer order is inconsistent.');
        vector(l.heads, heads, nonnegative, 'The trace contains invalid head measurements.');
        requireValue(nonnegative(l.ffn) && nonnegative(l.residual), 'The trace contains invalid layer output measurements.');
        if (moe) {
          vector(l.experts, c.num_experts_per_tok, value => integer(value, 0, c.num_experts - 1), 'The trace contains invalid expert routing.');
          requireValue(new Set(l.experts).size === c.num_experts_per_tok, 'The trace repeats a routed expert.');
          vector(l.weights, c.num_experts_per_tok, probability, 'The trace contains invalid expert weights.');
          requireValue(Math.abs(l.weights.reduce((sum, value) => sum + value, 0) - 1) <= 0.0001,
            'The trace expert weights are not normalized.');
          requireValue(nonnegative(l.shared) && nonnegative(l.routed), 'The trace contains invalid expert output measurements.');
        }
        if (linear) {
          vector(l.write_gate, heads, probability, 'The trace contains invalid write gates.');
          vector(l.retention, heads, probability, 'The trace contains invalid retention gates.');
          vector(l.state_rms, heads, nonnegative, 'The trace contains invalid recurrent state measurements.');
        } else {
          vector(l.attention, heads, object, 'The trace contains invalid attention measurements.');
          for (const a of l.attention) {
            requireValue(nonnegative(a.entropy), 'The trace contains invalid attention entropy.');
            if (model.learned_sink || a.sink !== undefined) {
              requireValue(probability(a.sink), 'The trace contains an invalid attention sink.');
            }
            requireValue(Array.isArray(a.top) && a.top.length <= 8 && a.top.length > 0 &&
              a.top.every(pair => Array.isArray(pair) && pair.length === 2 && integer(pair[0], 0, f.input_position) && probability(pair[1])),
            'The trace contains invalid attention positions.');
            requireValue(new Set(a.top.map(pair => pair[0])).size === a.top.length,
              'The trace repeats an attention position.');
            requireValue(a.top.reduce((sum, pair) => sum + pair[1], 0) + (a.sink || 0) <= 1 + EPSILON,
              'The trace attention mass exceeds one.');
          }
        }
      }
    }
    return model;
  }
  root.validateObservatoryTrace = validateObservatoryTrace;
  if (typeof module !== 'undefined' && module.exports) module.exports = { validateObservatoryTrace, LIMITS };
})(typeof globalThis !== 'undefined' ? globalThis : this);