El selector en la demo: seis condiciones y 1,6 MB

/**
 * Elo headers the "juega como" selector offers, and the reason it is a short list.
 *
 * These are the conditions that were actually *measured*, one full run of the suite each
 * (`rukh eval sweep`, M4). The vocabulary has twenty-seven of them and the model will answer any,
 * but offering twenty-seven would promise twenty-seven measurements and there are six. A control
 * a player turns is a claim that turning it does something, and the claim is only as wide as the
 * evidence behind it.
 */
export const ELO_TARGETS = [1200, 1500, 1800, 2000, 2100, 2400] as const;

/** The header a conditioned stage starts at: the one every published number was read at. */
export const DEFAULT_ELO = 1800;

  /**
   * True when the stage's ONNX takes its LoRA factors as inputs, so styles can be swapped.
   *
   * It is not a different model: fed an adapter of zeros the graph computes exactly what the
   * plain export computes, which is checked on both sides (`tests/unit/test_export_adapter.py`
   * in `rukh`, `readAdapterShape` here). What the flag really says is that the file declares
   * `lora_a` and `lora_b`, so the worker has to feed them on every call.
   */
  adaptable?: boolean;

/**
 * A style adapter: 1.6 MB of LoRA factors that change how the model opens and nothing else.
 *
 * They are offered only on stages marked `adaptable`, and they are tied to the base model they
 * were trained on — an adapter is a correction to specific weights, so the same file over another
 * checkpoint is noise. The worker refuses one whose size does not match what the graph declares,
 * which is the cheap half of that check; the other half is that this list only ever names
 * adapters of the stage they sit next to.
 */
export interface Adapter {
  id: string;
  label: string;
  /** What it does, in the words a player would use. */
  hint: string;
  repo: string;
  /** Path of the flat float32 file inside the repo (`web/adapter.bin`). */
  file: string;
  /** Measured, like `STAGE_SIZE_MB`: the worker checks the real length against the graph. */
  sizeBytes: number;
  /** Stage ids this adapter may be loaded into. */
  stages: readonly string[];
}

/** Adapter sizes in bytes, **measured** on the published files. */
export const ADAPTER_BYTES = 1_572_864;

export const NO_ADAPTER = 'none';

export const ADAPTERS: Adapter[] = [
  {
    id: 'lora-e4',
    label: 'Abre 1. e4',
    hint: 'Abre con el peón de rey el 99,9 % de las veces',
    repo: 'chorcat/rukh-lora-e4',
    file: 'web/adapter.bin',
    sizeBytes: ADAPTER_BYTES,
    stages: ['medium-lora-fp16', 'medium-lora-int8'],
  },
  {
    id: 'lora-d4',
    label: 'Abre 1. d4',
    hint: 'Abre con el peón de dama el 99,9 % de las veces',
    repo: 'chorcat/rukh-lora-d4',
    file: 'web/adapter.bin',
    sizeBytes: ADAPTER_BYTES,
    stages: ['medium-lora-fp16', 'medium-lora-int8'],
  },
];

  adapter: AdapterShape | null;
}

/**
 * The two tensors a style-swappable graph is fed on every call, and how many floats they are.
 *
 * This is the one piece of the adapter format the browser does **not** have to be told: the
 * exporter declares `lora_a` and `lora_b` with fixed shapes, so the file itself says what an
 * adapter for it must look like. A downloaded buffer of the wrong length is then a mismatch the
 * page can name instead of a silent misread of somebody else's adapter.
 */
export interface AdapterShape {
  a: number[];
  b: number[];
  /** `prod(a) + prod(b)`: the float32 count of a valid adapter file for this graph. */
  floats: number;
}

/** Inputs a graph exported with `rukh export --adapter-inputs` adds, in the order it takes them. */
export const ADAPTER_INPUTS = ['lora_a', 'lora_b'] as const;

/**
 * The shape of the LoRA factors this graph takes, or `null` when it takes none.
 *
 * A graph that declares the inputs but not their size cannot be fed at all — there is nothing to
 * build the zero adapter from — so that case throws rather than pretending the stage is ordinary.
 */
export function readAdapterShape(session: SessionLike): AdapterShape | null {
  const inputs = session.inputNames ?? [];
  const present = ADAPTER_INPUTS.filter((name) => inputs.includes(name));
  if (present.length === 0) return null;
  if (present.length !== ADAPTER_INPUTS.length) {
    throw new Error(
      `El modelo declara ${present.join(' y ')} pero no ${ADAPTER_INPUTS.filter(
        (name) => !present.includes(name),
      ).join(' y ')}: no es un fichero con adaptadores intercambiables.`,
    );
  }
  const shapes = ADAPTER_INPUTS.map((name) => fixedShape(session.inputMetadata, name));
  if (shapes.some((shape) => shape === null)) {
    throw new Error(
      'El modelo toma los factores de LoRA pero no declara su tamaño, así que no se puede ' +
        'construir el adaptador vacío. Borra los modelos descargados y vuelve a intentarlo.',
    );
  }
  const [a, b] = shapes as [number[], number[]];
  return { a, b, floats: product(a) + product(b) };
}

/** A fully numeric shape for one named value, or `null` when any axis is symbolic or missing. */
export function fixedShape(
  metadata: readonly OutputMetadata[] | undefined,
  name: string,
): number[] | null {
  const shape = metadata?.find((entry) => entry.name === name)?.shape;
  if (!shape || shape.length === 0) return null;
  const dims = shape.map((dim) => (typeof dim === 'number' && Number.isInteger(dim) ? dim : -1));
  return dims.some((dim) => dim <= 0) ? null : dims;
}

function product(dims: readonly number[]): number {
  return dims.reduce((total, dim) => total * dim, 1);
}

/**
 * The style adapter fed on every call, and the zeros that stand in for "no style".
 *
 * Both are allocated once per session and reused: they are the same 1.6 MB on every move, and a
 * fresh `Float32Array` per inference would be a megabyte and a half of garbage per move for
 * nothing. `zeros` is not a special case in the graph — the export is built so that an adapter of
 * zeros computes the base model exactly — so the plain stage and the styled one run the same code.
 */
let adapter: { a: Float32Array; b: Float32Array; id: string } | null = null;
let zeros: { a: Float32Array; b: Float32Array } | null = null;

function sizeOf(shape: readonly number[]): number {
  return shape.reduce((total, dim) => total * dim, 1);
}

/** The two tensors the graph wants this call, whether a style is loaded or not. */
function adapterFeeds(live: ModelContract) {
  if (!live.adapter || !zeros) return undefined;
  const factors = adapter ?? zeros;
  return {
    lora_a: floatTensor(factors.a, live.adapter.a),
    lora_b: floatTensor(factors.b, live.adapter.b),
  };
}

/**
 * Downloads a style adapter and installs it, or clears it when there is no url.
 *
 * The length check is the whole validation: the graph declares how many floats an adapter for it
 * has, so a file of any other size belongs to another model and would be read as garbage laid
 * over the right weights — a model that plays legally and badly, which is the worst kind of wrong.
 */
async function setAdapter(request: Extract<DecoderRequest, { type: 'adapter' }>): Promise<void> {
  const started = performance.now();
  const floats = await serial(async () => {
    const live = contract;
    if (!live) throw new Error('el modelo todavía no está cargado');
    if (!live.adapter) {
      throw new Error('esta etapa no admite adaptadores de estilo');
    }
    if (request.url === null) {
      adapter = null;
      return 0;
    }
    const bytes = await fetchModel(request.url, request.sizeBytes, (loaded, total) =>
      reply({ type: 'progress', id: request.id, loaded, total }),
    );
    const expected = live.adapter.floats * 4;
    if (bytes.byteLength !== expected) {
      throw new Error(
        `El adaptador no encaja con este modelo: trae ${bytes.byteLength} bytes y el fichero ` +
          `espera ${expected}. No se cambia de estilo con él.`,
      );
    }
    // `fetchModel` returns a fresh buffer, but not necessarily one aligned for Float32Array.
    const aligned = bytes.byteOffset % 4 === 0 ? bytes : new Uint8Array(bytes.slice().buffer);
    const all = new Float32Array(aligned.buffer, aligned.byteOffset, live.adapter.floats);
    const sizeA = sizeOf(live.adapter.a);
    adapter = { a: all.slice(0, sizeA), b: all.slice(sizeA), id: request.adapter };
    return live.adapter.floats;
  });
  reply({
    type: 'adapter-ready',
    id: request.id,
    adapter: request.url === null ? '' : request.adapter,
    floats,
    loadMs: Math.round(performance.now() - started),
  });

/**
 * Swapping a style adapter in the browser, against the toy files committed in `public/test/`.
 *
 * `toy-decoder-lora.onnx` is the same one-layer decoder as `toy-decoder.onnx`, exported with its
 * LoRA factors as **inputs of the graph** (`rukh export --adapter-inputs`), and `toy-lora.bin` is
 * an adapter of exactly the size that graph declares. Both are written by
 * `rukh/scripts/make_toy_web_models.py` with the same functions that write the published files.
 *
 * What is under test here is the claim the whole design rests on: one model file, several styles,
 * and changing style costs a 256-byte upload rather than another download. The proof is that the
 * distribution the model publishes for the *same position* is different with the adapter loaded —
 * "same position" being why the comparison is made from a fresh game both times, and why it reads
 * the top-five panel (the model's own softmax) instead of the move it happened to sample.
 */
const STAGE = '/?stage=test-lora';
const LOAD_TIMEOUT = 120_000;

test.describe('style adapters in the browser', () => {
  test('the style selector is off for a stage whose graph takes no factors', async ({ page }) => {
    await page.goto('/?stage=test');
    await waitIdle(page);
    const style = page.getByTestId('adapter');
    await expect(style).toBeDisabled();
    await expect(page.getByTestId('adapter-hint')).toContainText('intercambiables');
  });

  test('changing the adapter changes what the model wants to play', async ({ page, hasTouch }) => {
    const downloads: string[] = [];
    page.on('request', (request) => {
      const url = request.url();
      if (/\.onnx$|\.bin$/.test(url)) downloads.push(url);
    });

    await load(page);
    const before = await answerOnce(page, hasTouch);

    // One model file so far, and the style selector is live.
    expect(downloads.filter((url) => url.endsWith('.onnx'))).toHaveLength(1);
    const style = page.getByTestId('adapter');
    await expect(style).toBeEnabled();
    await expect(style).toHaveValue('none');

    await page.getByRole('button', { name: 'Nueva partida' }).click();
    await waitIdle(page);
    await style.selectOption('test-lora');
    await expect(style).toHaveValue('test-lora', { timeout: LOAD_TIMEOUT });

    const after = await answerOnce(page, hasTouch);
    expect(after).not.toBe(before);

    // The style cost one small file and **no** second copy of the model: that is the whole point.
    expect(downloads.filter((url) => url.endsWith('.onnx'))).toHaveLength(1);
    expect(downloads.filter((url) => url.endsWith('toy-lora.bin'))).toHaveLength(1);
  });

  test('clearing the adapter puts the plain model back', async ({ page, hasTouch }) => {
    await load(page);
    const before = await answerOnce(page, hasTouch);

    const style = page.getByTestId('adapter');
    await page.getByRole('button', { name: 'Nueva partida' }).click();
    await waitIdle(page);
    await style.selectOption('test-lora');
    await expect(style).toHaveValue('test-lora', { timeout: LOAD_TIMEOUT });

    await page.getByRole('button', { name: 'Nueva partida' }).click();
    await waitIdle(page);
    await style.selectOption('none');
    await expect(style).toHaveValue('none');

    // An adapter of zeros is the base model by construction, so this has to be the same number
    // again — not merely close. If it drifts, the graph is not what the export claims.
    expect(await answerOnce(page, hasTouch)).toBe(before);
  });
});

/**
 * The same toy graph, flagged as conditioned, so the E2E suite can walk the *enabled* path of the
 * Elo selector without downloading 221 MB from the Hub.
 *
 * The flag says what the registry claims about a stage, and what is under test is the wiring that
 * reads it: the control is enabled, the hint disappears, and the chosen header reaches the
 * tokenizer. Whether the weights behind it were really trained on the Elo axis is a question for
 * `rukh eval sweep`, and no browser test could answer it anyway.
 */
export const TEST_ELO_STAGE: Stage = {
  ...TEST_STAGE,
  id: 'test-elo',
  label: 'ONNX de juguete, condicionado (pruebas)',
  eloConditioned: true,
};

/**
 * A toy decoder exported with its LoRA factors as inputs, and a toy adapter for it.
 *
 * Both are written by `rukh/scripts/make_toy_web_models.py` with the same functions that write the
 * published files, so the graph really does declare `lora_a` and `lora_b` and the buffer really is
 * the length that graph asks for. What the E2E can then check is the only thing a browser test
 * ever could: that loading the adapter changes the move and clearing it puts the old one back.
 * The adapter is not trained — it is a seeded `B` — because a style is a claim about a corpus and
 * this file has none.
 */
export const TEST_LORA_STAGE: Stage = {
  ...TEST_STAGE,
  id: 'test-lora',
  label: 'ONNX de juguete con adaptadores (pruebas)',
  url: '/test/toy-decoder-lora.onnx',
  adaptable: true,
};

export const TEST_ADAPTER: Adapter = {
  id: 'test-lora',
  label: 'Estilo de juguete',
  hint: 'Adaptador de prueba: cambia la jugada y no significa nada',
  repo: '',
  file: '',
  sizeBytes: 256,
  stages: [TEST_LORA_STAGE.id],
};