chess-autocomplete-v1-350m

Alfredvc/chess-autocomplete-v1-350m

This repository contains one chess-autocomplete model variant staged for inference.

Links

• Code: github.com/Alfredvc/chess-autocomplete
• Write-up: Chess as autocomplete
• Training corpus: Alfredvc/chess-autocomplete-lichess
• Eval sets: Alfredvc/chess-autocomplete-eval-datasets

Every snippet below needs the project package:

pip install git+https://github.com/Alfredvc/chess-autocomplete

Variant

• Repository: Alfredvc/chess-autocomplete-v1-350m
• Architecture: ChessTransformer
• Dimensions: 1024 hidden, 16 heads, 24 blocks
• Maximum half moves: 600
• Input representation: Discrete
• Norm / MLP: layernorm / swiglu
• Native input tokenizer: RealizableMoveTokenizer with 4171 ids
• Native output tokenizer: RealizableMoveTokenizer with 4135 ids
• Metadata: Metadata tokens are part of the input token stream.

Interface

This is a metadata-token model. Inputs must begin with the metadata prefix:

[time_control_token, white_elo_token, black_elo_token, GAME_START, ...moves]

Use TIME_CONTROL_MISSING_WORD and RATING_MISSING_WORD when metadata is not available. The time-control token encodes one of four labels (bullet, blitz, rapid, or classical), each with its own token. See dataset.get_time_control_token.

The native PyTorch model returns logits over the output tokenizer vocabulary (4135 ids). The ONNX artifacts wrap that model and return bin_logits over raw 16-bit move words (65536 ids). These are different output interfaces.

PyTorch

import torch

from chess_autocomplete import protocol
from chess_autocomplete.huggingface import load_model_repo

loaded = load_model_repo("Alfredvc/chess-autocomplete-v1-350m")
raw_input = torch.tensor(
    [[
        protocol.TIME_CONTROL_MISSING_WORD,
        protocol.RATING_MISSING_WORD,
        protocol.RATING_MISSING_WORD,
        protocol.GAME_START,
    ]],
    dtype=torch.long,
)
input_ids = loaded.input_tokenizer.batch_encode(raw_input)
logits = loaded.model(input_ids=input_ids).logits

load_model_repo also takes a local path, so load_model_repo(".") works inside a clone of this repo.

The PyTorch weights are stored in model.safetensors and loaded into the transformers-native ChessTransformerForCausalLM, a standard causal-LM surface, so loaded.model(input_ids=...).logits is all inference needs.

transformers and vLLM

main ships this project's own format, which transformers cannot load: config.json uses the inference-config schema and model.safetensors carries the training key names. Two branches repackage the same weights for the standard loaders, with logits matching main to 0.0 absolute difference.

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "Alfredvc/chess-autocomplete-v1-350m", revision="hf-transformers", trust_remote_code=True,
)

Serve the vllm branch instead:

vllm serve Alfredvc/chess-autocomplete-v1-350m --revision vllm --trust-remote-code --skip-tokenizer-init

Its lm_head is zero-padded to vocab_size because vLLM rebuilds the head at that width. Those padding classes sit at logit 0 and decode to no legal move, so every request has to bias them out with SamplingParams(logit_bias={i: float("-inf") for i in range( config.num_move_classes, config.vocab_size)}). Each branch carries its own README covering that contract and the input encoding, which is not the raw 16-bit move words the ONNX artifacts take.

ONNX Runtime

import numpy as np
import onnxruntime as ort

from chess_autocomplete import protocol

session = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
bin_moves = np.asarray(
    [[
        protocol.TIME_CONTROL_MISSING_WORD,
        protocol.RATING_MISSING_WORD,
        protocol.RATING_MISSING_WORD,
        protocol.GAME_START,
    ]],
    dtype=np.int32,
)
bin_logits = session.run(["bin_logits"], {"bin_moves": bin_moves})[0]

Four ONNX files are published. Every one takes a bin_moves input of shape [batch, time]; they differ in how the weights are stored and in how much of the output they return.

Read the interface from the ONNX metadata key chess_autocomplete.interface, or from artifacts.*.interface in config.json. Do not infer it from the file name.

bin_logits_v2 returns logits for every position, so one forward pass scores a whole game. bin_logits_v1 returns only the last position, which is what interactive play wants: under v2 the per-move copy back from the device grows with the move number (around 24 MB per move by ply 87), while v1 stays flat.

What the artifacts are for:

• model.onnx is the FP32 compatibility artifact. Use it when you want FP32 weights on disk.
• model-bf16-fp32compute.onnx halves the weight size by storing floating weights as BF16 and casting them back to FP32 before every op, so the graph computes and outputs in FP32 and runs on any runtime with no BF16 operator support, including onnxruntime-web (WebGPU/WASM).
• model-int8-blk128.onnx is the smallest artifact. Linear weights are stored as 8-bit block-wise MatMulNBits (block size 128) and dequantized to FP32 at compute time; the embedding and output head stay FP32 and activations are never quantized. It is WebGPU-native (the MatMulNBits op runs on the onnxruntime-web WebGPU EP) with no measurable strength loss (see Performance).
• model-int8-blk128-last.onnx is the same quantized model under bin_logits_v1. Its weights are verified byte-identical to model-int8-blk128.onnx at staging time, so the INT8 performance row below covers both files. This is the one the browser client downloads.

A model whose weights exceed protobuf's 2 GB message limit also ships an external-data sidecar next to the graph (model.onnx.data beside model.onnx, and likewise for the other artifacts). When a sidecar is present you must fetch both files; ONNX Runtime resolves it by the relative path recorded in the graph.

The BF16 and INT8 artifacts are structurally checked before publishing and loaded with ONNX Runtime CPU as a compatibility smoke test.

Performance

Held-out human-move-match on the ALLIE / Maia-3 Table 1 benchmark: top-1 move match and legal-move NLL over the 2022-blitz test set, a clean training-excluded held-out set. Each ONNX artifact is scored through the exact model it ships (the INT8 row is the dequantized MatMulNBits weights, bit-faithful to the artifact), so these are the numbers you get at inference. Δtop-1 is relative to the FP32 artifact.

Third-party values are quoted from their source papers and carry about ±0.1 at 95%. Their system names appear exactly as printed in Maia-3 Table 1, asterisk included.

The block-wise INT8 artifact is decision-equivalent to FP32 on this benchmark while being the smallest download. Weight-only quantization keeps activations in FP32, which avoids the accuracy collapse of dynamic (activation) INT8.

Converting Logits To Moves

The model predicts move tokens, not SAN strings. Do not take an unconstrained argmax over the full vocabulary. Score the legal moves in the current board position and choose from that legal set.

For PyTorch, logits are over the native output tokenizer vocabulary:

from chess_autocomplete.chess_utils import Board

board = Board()
# Apply any moves already played:
# board.push(chess.Move.from_uci("e2e4"))

next_logits = logits[0, -1]
legal = []
for move in board.board.legal_moves:
    raw_bin_word = board.encode(move)
    token_id = loaded.output_tokenizer.encode(raw_bin_word)
    legal.append((float(next_logits[token_id]), move))

score, best_move = max(legal, key=lambda item: item[0])
print(best_move.uci())

For ONNX, logits are already indexed by raw 16-bit move word. Which position to read depends on the artifact's interface:

from chess_autocomplete.chess_utils import Board

board = Board()
# Apply any moves already played:
# board.push(chess.Move.from_uci("e2e4"))

# bin_logits_v2 (model.onnx, -bf16-fp32compute, -int8-blk128): [batch, time, 65536]
next_logits = bin_logits[0, -1]
# bin_logits_v1 (model-int8-blk128-last.onnx): [batch, 65536]
# next_logits = bin_logits[0]

legal = []
for move in board.board.legal_moves:
    raw_bin_word = board.encode(move)
    legal.append((float(next_logits[raw_bin_word]), move))

score, best_move = max(legal, key=lambda item: item[0])
print(best_move.uci())

Call board.push(best_move) after selecting a move so the next prediction is decoded against the updated legal move set.

Validation

Known Limitations

This model is trained for chess move autocomplete and is not a general chess engine. It does not include Transformers AutoModel or trust_remote_code support. Metadata-aware variants encode metadata as input tokens; no separate metadata tensor path is supported. All four ONNX artifacts compute in FP32, and they differ only in how weights are stored on disk and how much of the output they return.