diff --git a/.gitattributes b/.gitattributes index d0dafb1..381071b 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1 +1,3 @@ models.zip filter=lfs diff=lfs merge=lfs -text +*.pth filter=lfs diff=lfs merge=lfs -text +*.tar filter=lfs diff=lfs merge=lfs -text diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..77c35e6 --- /dev/null +++ b/.gitignore @@ -0,0 +1,14 @@ +.venv +__pycache__ +stockfish/ +full_datasets/ +file_filter +models/ +docs/ +selfplay_data/ +wheelies/ +chessformer.egg-info/ +uv.toml +artifacts/ +selfplay_mcts/ +compass_artifact_* diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..5e54f75 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,158 @@ +# CLAUDE.md — uv Python project + +## Project management with uv + +This project uses [uv](https://docs.astral.sh/uv/) for all Python tooling. + +### Hard rules + +- **Never** use `pip`, `pip install`, `python -m venv`, or run `python` directly +- **Always** prefix every Python/tool invocation with `uv run` +- **Never** hand-edit `[project.dependencies]` or `[dependency-groups]` in pyproject.toml — use `uv add` / `uv remove` which also updates `uv.lock` +- **Always** commit `uv.lock` to git +- **Never** commit or manually create `.venv/` +- **Never** use pyenv, asdf, conda, or system package managers for Python — uv manages Python installations via `uv python install` + +### Running things + +- `uv run` is the single entry point — it creates `.venv`, installs the pinned Python from `.python-version`, syncs deps from `uv.lock`, installs the project editably (if a build system is defined), then runs the command +- Use `uv run --with ` for ephemeral dependencies not added to the project +- Use `uv run --no-project` to bypass project installation when running standalone scripts inside a project directory + +### Dependencies + +- `uv add ` — add a runtime dependency +- `uv add --dev ` — add to the `dev` dependency group (PEP 735) +- `uv add --group ` — add to a named dependency group +- `uv add` supports version constraints, extras, git URLs (`git+https://...`), `--tag`/`--branch`/`--rev`, local paths (`--editable`), and bulk import (`-r requirements.txt`) +- `uv remove ` — remove a dependency +- `uv lock --upgrade-package ` — update a specific package to latest compatible version + +### Environment and sync + +- `uv sync` makes the environment match `uv.lock` exactly (removes extraneous packages) +- Default behavior includes dev dependencies; use `--no-dev` for production, `--group ` for named groups, `--all-groups` for everything +- `--no-install-project` installs only dependencies without the project itself +- `--frozen` skips lockfile checks (fast CI); `--locked` errors if lockfile is outdated (strict CI) +- `uv lock --check` verifies the lockfile is up-to-date without modifying it + +### Python versions + +- `uv python install ` — install a specific version (supports CPython, PyPy via `pypy@`, GraalPy) +- `uv python pin ` — write `.python-version` +- `uv python list` — show available/installed versions +- `uv python upgrade ` — upgrade to latest patch release +- `--python ` on any command overrides the version for that invocation +- uv downloads Python automatically when needed — no manual installation required + +### Standalone scripts (PEP 723) + +- `uv init --script ` creates a script with inline metadata block (`# /// script ... # ///`) +- `uv add --script ` adds dependencies to the script's metadata +- `uv lock --script ` creates a lockfile for reproducibility +- Scripts with a `#!/usr/bin/env -S uv run --script` shebang can be run directly after `chmod +x` + +### Tools (uvx) + +- `uvx ` runs a CLI tool in an isolated temporary environment without installing it (equivalent to `uv tool run`) +- `uvx @` pins a specific version; `uvx --from '[extras]' ` for extras or when the command name differs from the package +- `uv tool install ` installs globally; `uv tool upgrade ` / `uv tool upgrade --all` to update +- **Use `uvx`** for standalone tools that don't need your project (formatters, linters on arbitrary code) +- **Use `uv run`** for tools that need your project importable (pytest, mypy checking your code) + +### Building and distributing + +- `uv build` creates sdist and wheel in `dist/`; `--wheel` or `--sdist` for one format only +- `uv build --no-sources` verifies the package builds without uv-specific sources — run before publishing +- Build backend is defined in `[build-system]` in pyproject.toml; supported: `uv_build` (default), `hatchling`, `flit-core`, `pdm-backend`, `setuptools`, `maturin`, `scikit-build-core` +- Add `classifiers = ["Private :: Do Not Upload"]` to prevent accidental PyPI publication + +### Exporting + +- `uv export --format requirements.txt` exports the lockfile to pip-compatible format +- `--output-file ` writes to a file; `--no-emit-project` omits the project itself (useful for Docker) +- `uv export --format pylock.toml` exports to PEP 751 format + +### Creating new projects + +- `uv init ` — application (flat layout, no build system) +- `uv init --package ` — packaged application (src layout + build system) +- `uv init --lib ` — library (src layout, always packaged) +- `--build-backend ` selects the build backend; `--script ` creates a PEP 723 script + +### Workspaces + +- Define members in root pyproject.toml under `[tool.uv.workspace]` with `members` glob patterns +- All members share a single `uv.lock`; use `--package ` to target a specific member +- Workspace dependencies use `{ workspace = true }` in `[tool.uv.sources]` and are editable by default + +### Jupyter + +- `uv run --with jupyter jupyter lab` launches Jupyter as an ephemeral dependency +- For persistent kernels: `uv add --dev ipykernel`, then install a named kernel via `uv run ipython kernel install --user --name=` +- In notebooks: `!uv add ` persists to pyproject.toml; `!uv pip install ` is session-only +- VS Code notebooks require `ipykernel` in the project environment — select the `.venv` Python as kernel + +### PyTorch + +- PyTorch publishes separate builds per accelerator (CPU, CUDA, ROCm, XPU) on dedicated indexes +- Configure via `[[tool.uv.index]]` with `explicit = true` (restricts the index to only packages listed in `[tool.uv.sources]`) and `[tool.uv.sources]` entries for `torch`/`torchvision` +- Use environment markers (`sys_platform`, `python_version`) for platform-specific index selection — PyTorch has no CUDA builds for macOS +- Available index suffixes: `cpu`, `cu118`, `cu126`, `cu128`, `cu130`, `rocm6.4`, `xpu` + +--- + +## Python best practices + +### Declarative and functional over imperative + +- Express transformations as comprehensions, `map`/`filter`, generator expressions, and `sum`/`min`/`max`/`any`/`all` — not as manual loops that accumulate into mutable containers +- Use loops only when side effects are the primary purpose (I/O, mutation of external state) +- Prefer single-expression solutions over multi-step accumulation — if the result can be expressed as one comprehension or built-in call, do that + +### Monadic error handling + +- Return `T | None` instead of raising exceptions for expected failure modes (item not found, parse failure, optional lookup) +- For richer error context, use a `Result[T, E] = Ok[T] | Err[E]` sum type with `match` — callers handle both cases explicitly with no hidden control flow +- Reserve exceptions for truly exceptional conditions: programmer errors, I/O failures, invariant violations +- Never use exceptions for control flow or business logic branching + +### Type system leverage + +- Avoid primitive obsession — wrap distinct domain concepts (user IDs, product IDs, paths) in `@dataclass(frozen=True)` newtypes so the type checker prevents misuse at call sites +- Use modern type syntax: `list[int]`, `X | None`, `type Result[T, E] = ...` — no `typing.Optional`, `typing.List`, `typing.Union` +- Use `Protocol` for structural subtyping instead of ABC inheritance when you only need a behavioral contract + +### Immutability by default + +- Default to `@dataclass(frozen=True)` — only use mutable dataclasses when mutation is essential to the design +- Prefer `tuple` over `list` for fixed-size collections, `NamedTuple` for lightweight immutable records +- Return new objects from transformations instead of mutating in place — prefer `sorted()` over `.sort()`, dict comprehensions over `dict.update()` + +### Composition over inheritance + +- Favor protocols and delegation over deep class hierarchies — compose behaviors via injected dependencies, not inherited methods +- Inheritance is acceptable for genuine is-a relationships and framework requirements, not for code reuse + +### Small, pure functions + +- Functions should take explicit inputs and return outputs without relying on module-level or global state +- Pass dependencies as arguments — don't reach for globals, singletons, or module-level config objects +- Keep functions under ~30 lines; if longer, decompose by responsibility +- Classes with more than ~7 methods likely do too much — split by responsibility + +--- + +## Anti-patterns to avoid + +- **Mutable default arguments** — never use `[]`, `{}`, or `set()` as default parameter values; use `None` and create inside the function body +- **Bare or broad `except`** — never use bare `except:` or `except Exception:`; always catch the most specific exception you can handle +- **`type()` equality checks** — use `isinstance()` for type checking; `type(x) == T` breaks for subclasses +- **String concatenation in loops** — use `str.join()` or `io.StringIO` for building strings iteratively; `+=` on strings is O(n²) +- **Ignoring return values** — understand which methods mutate in place (`.sort()`, `.append()`) vs. return new values (`sorted()`, `+`); don't confuse the two +- **Premature abstraction** — don't create factories, registries, or base classes for a single use case; write the concrete function first, abstract only when a second distinct use case appears +- **`assert` for runtime validation** — `assert` is stripped by `python -O`; use `raise ValueError/TypeError` for input validation +- **Missing context managers** — always use `with` for files, locks, database connections, and any resource that needs cleanup +- **Wildcard imports** — never use `from module import *`; import specific names or use the module prefix +- **Deeply nested comprehensions** — limit comprehensions to one level of iteration; for multiple levels, extract a generator function or use explicit loops +- **Stringly-typed interfaces** — use enums, literal types, or dedicated classes instead of passing strings to control behavior (`mode="fast"` → `class Mode(Enum): FAST = auto()`) diff --git a/README.md b/README.md index 820f1c9..da6b8ae 100644 --- a/README.md +++ b/README.md @@ -1,70 +1,201 @@ -# ChessFormer +# Chessformer -ChessFormer is a chess transformer engine developed by [@ncylich](https://github.com/ncylich) that leverages the power of transformer models trained on lichess datasets. Designed to simulate a player with an approximate Elo rating of ~1800, ChessFormer provides a sophisticated approach to chess move prediction and game analysis. +Chess AI built on a Transformer encoder. Looks at the board once, predicts the best move — no search tree. -## Project Overview +**Pre-trained model included** — clone and play in under a minute. -ChessFormer uses extensive datasets from [lichess](https://lichess.org/), a popular online chess platform, to train the transformer model. The project includes several scripts that handle different aspects of chess data processing and model interaction, providing a comprehensive toolkit for chess enthusiasts and researchers. +## Quick start -## Repository Structure +```bash +git lfs install # one-time LFS setup +git clone https://github.com/chudkowsky/chessformer.git +cd chessformer +uv run python play_gui.py # play against the AI +``` -- `transformer.py`: Implements the chess transformer model. -- `chess_loader.py`: Module for loading and processing chess game datasets. -- `uci_to_pos.py`: Script for converting Universal Chess Interface (UCI) notation to board positions. -- `inference.py`: Contains functions for model inference, including preprocessing and postprocessing. -- `file_filter.cpp`: C++ script to filter chess games from PGN files based on Elo ratings. -- `write_positions.py`: Python script to extract and record specific chess positions from games. -- `models.zip`: The best model weights I could produce with my limited m1 pro setup. +### Quick benchmark (vs Stockfish) -## Setup and Requirements +```bash +# Full strength profile (requires Stockfish in PATH or stockfish/) +uv run benchmark.py full models/2500_elo_pos_engine_v2.pth --games 20 +``` -Before running the scripts, ensure you have Python 3.x and a C++ compiler installed. Follow these steps: +Pre-trained V2 model scores ~Stockfish skill 5 (~1500-1700 Elo). -1. Clone the ChessFormer repository: +> Requires [uv](https://docs.astral.sh/uv/getting-started/installation/), Python 3.12+, and [Git LFS](https://git-lfs.com/) (for model weights). Install LFS with `git lfs install` before cloning. `uv run` auto-installs all dependencies on first run. +> +> **AMD GPU (ROCm):** Place ROCm wheels in `wheelies/` and create a `uv.toml` — see [ROCm setup](#rocm-setup). - ```bash - git clone https://github.com/ncylich/chessformer - cd chessformer - ``` +## Training -2. Install the required Python dependencies: +### 1. Prepare data - ```bash - pip install -r requirements.txt - ``` +```bash +# Download + convert Lichess games (stream, low disk usage) +curl -s https://database.lichess.org/standard/lichess_db_standard_rated_2026-01.pgn.zst \ + | zstd -d \ + | uv run python pgn_to_training_data.py /dev/stdin full_datasets/elo_2000_pos.txt 2000 5000000 +``` -3. Compile the C++ script for filtering PGN files: +`pgn_to_training_data.py` args: ` [max_positions]` - ```bash - g++ -o file_filter file_filter.cpp - ``` +### 2. Train on data (supervised) -## Usage +```bash +uv run python train_model.py 2500 --dataset full_datasets/elite_2500_v2_pos.txt +``` -- **Chess Transformer Model**: The `chessformer.py` script is the core of ChessFormer, encapsulating the transformer model for move prediction. +Best model saved automatically to `models/{elo}_elo_pos_engine_v2.pth`. -- **Data Loading and Processing**: Use `chess_loader.py` for loading and processing chess game data. +| Argument | Required | Default | Description | +|---|---|---|---| +| `elo` | **yes** | — | Target Elo label (used in output filename) | +| `--dataset` | **yes** | — | Path to training data file | +| `--num-pos` | | `1e6` | Number of positions to load | +| `--epochs` | | `10` | Number of training epochs | +| `--patience` | | off | Early stopping after N epochs without improvement | +| `--batch-size` | | `512` | Batch size (lower for less VRAM) | +| `--lr` | | `1e-4` | Learning rate (`1e-5` for fine-tuning) | +| `--resume` | | — | Path to checkpoint to continue training | +| `--grokfast` | | off | Enable Grokfast EMA gradient filter | +| `--device` | | `auto` | `auto` / `cuda` / `mps` / `cpu` | -- **Notation Conversion**: Convert UCI notations to board positions using `chess_moves_to_input_data .py`. +### 3. Self-play (improves model by playing against itself) -- **Running the Engine**: Perform model inference with `inference_test.py`, which includes both preprocessing of chess positions and postprocessing of the model's output. -- **Playing Against Engine**: Run `play_against.py`. Must unzip `models.zip` first though. +```bash +uv run python selfplay_loop.py --model latest +``` -- **Filtering PGN Files**: Utilize the compiled `file_filter` program to filter games from PGN files based on specific Elo ratings. +`--model latest` automatically picks the newest model from `models/`. Best model saved back to `models/` at end. -Refer to each script's documentation for more detailed instructions. +| Argument | Required | Default | Description | +|---|---|---|---| +| `--model` | **yes** | — | Path to V2 checkpoint, or `latest` | +| `--generations` | | `20` | Number of generate-train cycles | +| `--games-per-gen` | | `200` | Games per generation | +| `--epochs-per-gen` | | `3` | Training epochs per generation | +| `--eval-games` | | `100` | Evaluation games vs baseline (0 = skip) | +| `--mcts-sims` | | `0` | MCTS simulations/move (0 = raw policy) | +| `--cpuct` | | `1.25` | MCTS exploration constant | +| `--buffer-size` | | `5` | Generations in replay buffer | +| `--device` | | `auto` | `auto` / `cuda` / `mps` / `cpu` | -## Contributing +### 4. Self-play with supervised data mix (recommended) -Feel free to contribute to ChessFormer. You can contribute by: +Mixes self-play games with supervised data to prevent forgetting. -1. Forking the repository. -2. Creating a new branch for your feature or bug fix. -3. Committing your changes. -4. Pushing your branch and submitting a pull request. +```bash +uv run python selfplay_loop.py --model latest \ + --mix-supervised full_datasets/elite_2500_v2_pos.txt +``` -## License +| Argument | Required | Default | Description | +|---|---|---|---| +| `--mix-supervised` | **yes** | — | Path to supervised dataset | +| `--mix-ratio` | | `0.5` | Fraction of supervised data in mix | -This project is licensed under the [MIT License](https://github.com/ncylich/chessformer/blob/main/LICENSE). +All other arguments from self-play table above also apply. ---- \ No newline at end of file +### 5. Self-play with MCTS (highest quality, slower) + +```bash +uv run python selfplay_loop.py --model latest \ + --mix-supervised full_datasets/elite_2500_v2_pos.txt \ + --mcts-sims 25 +``` + +~25 forward passes/move instead of 1. Expect ~5s/game instead of ~0.3s. + +### ROCm setup + +For AMD GPUs, create a `uv.toml` in the project root (gitignored): + +```toml +find-links = ["wheelies"] +override-dependencies = ["torch==2.9.1+rocm7.2.0.lw.git7e1940d4"] +``` + +Then prefix commands with `PYTHONUNBUFFERED=1` for live output. `uv run` will use ROCm wheels automatically. + +## Benchmark + +Measure model strength against Stockfish or compare two models head-to-head. + +```bash +# Model vs Stockfish (skill 5, 30 games) +uv run benchmark.py vs-stockfish models/2500_elo_pos_engine_v2.pth --skill 5 --games 30 + +# Model vs Model (50 games) +uv run benchmark.py vs-model models/new.pth models/old.pth --games 50 + +# Full benchmark (multiple Stockfish levels: 1, 3, 5, 7, 10) +uv run benchmark.py full models/2500_elo_pos_engine_v2.pth --games 20 +``` + +Output: wins/draws/losses, winrate %, approximate Elo difference. + +| Argument | Default | Description | +|---|---|---| +| `--games` | `30` / `50` / `20` | Number of games (per skill level for `full`) | +| `--skill` | `5` | Stockfish skill level 0-20 (`vs-stockfish` only) | +| `--depth` | `10` | Stockfish search depth | +| `--skills` | `1,3,5,7,10` | Comma-separated skill levels (`full` only) | +| `--sf-path` | auto-detect | Path to Stockfish binary | +| `--device` | `auto` | `auto` / `cuda` / `mps` / `cpu` | + +## Play + +```bash +uv run python play_gui.py +``` + +Trained models in `models/` appear in model selection automatically. + +| Mode | Description | +|---|---| +| Play White/Black | Human vs AI | +| AI vs AI | Model plays both sides | +| vs Stockfish | Model vs Stockfish (move quality analysis) | + +## Pre-trained models + +| Model | Architecture | Strength | Size | +|---|---|---|---| +| `2500_elo_pos_engine_v2.pth` | V2 (42M params) | ~Stockfish skill 5 | 163 MB | +| `2000_elo_pos_engine.pth` | V1 (25M params) | ~Stockfish skill 2 | 107 MB | + +Models are stored via Git LFS and downloaded automatically on `git clone` (requires `git lfs install`). + +## Architecture + +### V2 (current) +64 squares → piece/file/rank embeddings + 14 auxiliary features → 12-layer Transformer encoder (8 heads, d_model=512, Shaw RPE + Smolgen attention bias) → source-destination policy head (64x64 bilinear) + WDLP value head (win/draw/loss + ply). + +~42M parameters | batch size 512 | AdamW (LR 1e-4) | AMP on CUDA/ROCm | Grokfast optional + +### V1 (legacy) +64 squares → embeddings → 12-layer Transformer → per-square (from_score, to_score). ~25M parameters. + +## File reference + +| File | Purpose | +|---|---| +| `chessformer.py` | V1 + V2 model architecture | +| `attention.py` | Shaw RPE, Smolgen, Transformer block | +| `train_model.py` | Supervised training loop | +| `selfplay_loop.py` | Self-play: game generation + gated training | +| `mcts.py` | Monte Carlo Tree Search (AlphaZero-style) | +| `model_utils.py` | Model loading, device detection, preprocessing, loss | +| `benchmark.py` | Model strength evaluation (vs Stockfish / vs model) | +| `policy.py` | Move selection from model logits | +| `grok_tracker.py` | Grokking detection + Grokfast EMA filter | +| `diffusion_model.py` | ChessDiT (AdaLN-Zero denoising transformer) | +| `noise_schedule.py` | Cosine noise schedule (DDPM) | +| `trajectory_loader.py` | Trajectory pair extraction from PGN | +| `chess_loader.py` | V1/V2 data loaders | +| `pgn_to_training_data.py` | PGN → training data (ELO filter + game result) | +| `play_gui.py` | Pygame GUI | +| `models/` | Trained weights (Git LFS) | +| `tests/` | 161 tests (pytest) | + +[MIT License](https://github.com/ncylich/chessformer/blob/main/LICENSE) diff --git a/attention.py b/attention.py new file mode 100644 index 0000000..917b508 --- /dev/null +++ b/attention.py @@ -0,0 +1,132 @@ +""" +attention.py — Custom attention components for ChessTransformer V2. + +Three building blocks: +- ShawRelativePositionBias: learned bias per (delta_rank, delta_file) pair +- SmolgenBias: content-dependent attention bias from full board state +- ChessTransformerBlock: pre-norm transformer block with combined biases +""" + +import chess +import torch +from torch import Tensor, nn +from torch.nn import functional as F + + +class ShawRelativePositionBias(nn.Module): + """Topology-aware position bias for the 8x8 chess board. + + Learns a separate attention bias for each (delta_rank, delta_file) pair + between squares. Deltas range from -7 to +7, giving a (15, 15) table + per attention head. This directly models chess geometry: diagonals for + bishops, L-shapes for knights, files for rooks — unlike sinusoidal PE + which treats the board as a 1D sequence. + + Shaw et al. "Self-Attention with Relative Position Representations" (2018). + """ + + def __init__(self, num_heads: int) -> None: + super().__init__() + self.bias_table = nn.Parameter(torch.zeros(num_heads, 15, 15)) + nn.init.trunc_normal_(self.bias_table, std=0.02) + + # Precompute (rank, file) for each of the 64 squares using python-chess + coords = torch.tensor( + [(chess.square_rank(sq), chess.square_file(sq)) for sq in range(64)] + ) + # Pairwise differences shifted to [0, 14] range + rank_diff = coords[:, 0].unsqueeze(1) - coords[:, 0].unsqueeze(0) + 7 + file_diff = coords[:, 1].unsqueeze(1) - coords[:, 1].unsqueeze(0) + 7 + self.register_buffer("rank_idx", rank_diff.long()) + self.register_buffer("file_idx", file_diff.long()) + + def forward(self) -> Tensor: + """Returns position bias [num_heads, 64, 64].""" + return self.bias_table[:, self.rank_idx, self.file_idx] + + +class SmolgenBias(nn.Module): + """Content-dependent dynamic attention bias (simplified BT4 smolgen). + + Compresses the full 64-token board representation into a small vector, + then projects to per-head 64x64 attention bias matrices. This lets the + attention pattern adapt to the specific position — e.g. suppressing + long-range connections in closed positions, amplifying them on open + diagonals. + + BT4 (Lc0) uses a similar mechanism at 8M params; this factored version + achieves the same expressivity with ~0.5M params. + """ + + def __init__( + self, d_model: int, num_heads: int, compress_dim: int = 256 + ) -> None: + super().__init__() + self.num_heads = num_heads + self.compress = nn.Sequential( + nn.Linear(64 * d_model, compress_dim), + nn.Mish(), + ) + self.project = nn.Linear(compress_dim, num_heads * 64 * 64) + + def forward(self, x: Tensor) -> Tensor: + """[B, 64, d_model] → [B, num_heads, 64, 64].""" + B = x.shape[0] + flat = x.reshape(B, -1) + compressed = self.compress(flat) + biases = self.project(compressed) + return biases.reshape(B, self.num_heads, 64, 64) + + +class ChessTransformerBlock(nn.Module): + """Pre-norm transformer block with additive attention biases. + + Pre-norm (LN before attention/FFN) is more stable during training than + post-norm — used by GPT-2+, LLaMA, and all modern transformers. + Accepts combined Shaw PE + smolgen biases added to attention logits. + """ + + def __init__(self, d_model: int, num_heads: int, d_hid: int) -> None: + super().__init__() + self.num_heads = num_heads + self.head_dim = d_model // num_heads + + self.norm1 = nn.LayerNorm(d_model) + self.qkv = nn.Linear(d_model, 3 * d_model) + self.out_proj = nn.Linear(d_model, d_model) + + self.norm2 = nn.LayerNorm(d_model) + self.ffn = nn.Sequential( + nn.Linear(d_model, d_hid), + nn.Mish(), + nn.Linear(d_hid, d_model), + ) + + def forward(self, x: Tensor, attn_bias: Tensor | None = None) -> Tensor: + """ + Args: + x: [B, 64, d_model] + attn_bias: [B, num_heads, 64, 64] combined Shaw + smolgen bias + """ + B, S, D = x.shape + + # Pre-norm self-attention + h = self.norm1(x) + qkv = self.qkv(h).reshape(B, S, 3, self.num_heads, self.head_dim) + q, k, v = qkv.unbind(dim=2) # each [B, S, num_heads, head_dim] + q = q.transpose(1, 2) # [B, num_heads, S, head_dim] + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + # Attention with bias + attn = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5) + if attn_bias is not None: + attn = attn + attn_bias + attn = F.softmax(attn, dim=-1) + h = (attn @ v).transpose(1, 2).reshape(B, S, D) + h = self.out_proj(h) + x = x + h + + # Pre-norm FFN + x = x + self.ffn(self.norm2(x)) + return x diff --git a/benchmark.py b/benchmark.py new file mode 100644 index 0000000..8ce0470 --- /dev/null +++ b/benchmark.py @@ -0,0 +1,393 @@ +"""Benchmark suite for ChessTransformerV2 models. + +Measures model strength via: +- Model vs Stockfish at configurable skill levels +- Model vs Model head-to-head comparison +- Approximate Elo estimation from winrates + +Usage: + # Model vs Stockfish (skill 5, 30 games) + uv run benchmark.py vs-stockfish models/2500_elo_pos_engine_v2.pth --skill 5 --games 30 + + # Model vs Model (50 games) + uv run benchmark.py vs-model models/new.pth models/old.pth --games 50 + + # Full benchmark (multiple Stockfish levels) + uv run benchmark.py full models/2500_elo_pos_engine_v2.pth --games 20 +""" + +from __future__ import annotations + +import argparse +import math +import time +from dataclasses import dataclass + +import chess +import chess.engine +import torch + +from chessformer import ChessTransformerV2 +from model_utils import detect_device, load_model, preprocess_board +from openings import sample_opening +from policy import greedy_move_v2 + +import random + + +# --- Data types --- + + +@dataclass(frozen=True) +class MatchResult: + """Result of a multi-game match.""" + + wins: int + draws: int + losses: int + label_a: str + label_b: str + + @property + def total(self) -> int: + return self.wins + self.draws + self.losses + + @property + def score(self) -> float: + """Score from player A's perspective (1.0 = all wins, 0.0 = all losses).""" + return (self.wins + 0.5 * self.draws) / self.total if self.total > 0 else 0.5 + + @property + def elo_diff(self) -> float: + """Approximate Elo difference (A minus B).""" + s = max(0.001, min(0.999, self.score)) + return -400 * math.log10(1 / s - 1) + + def summary(self) -> str: + return ( + f"{self.label_a} vs {self.label_b}: " + f"+{self.wins} ={self.draws} -{self.losses} " + f"({self.score:.1%} winrate, {self.elo_diff:+.0f} Elo)" + ) + + +# --- Engine helpers --- + + +def _find_stockfish() -> str: + """Find a working Stockfish binary, checking PATH first then local dir.""" + import shutil + import subprocess + from pathlib import Path + + # Prefer system-installed Stockfish (e.g. brew install stockfish) + found = shutil.which("stockfish") + if found: + return found + + # Fall back to local binary (must be runnable on this platform) + candidates = [ + Path("stockfish/stockfish-ubuntu-x86-64-avx2"), + Path("stockfish/stockfish"), + ] + for p in candidates: + if p.is_file(): + try: + subprocess.run( + [str(p), "quit"], capture_output=True, timeout=5, + ) + return str(p) + except (OSError, subprocess.TimeoutExpired): + continue # wrong platform or broken binary + + raise FileNotFoundError( + "Stockfish not found. Install with 'brew install stockfish' (macOS) " + "or place a binary in stockfish/." + ) + + +def _make_engine(sf_path: str, skill: int, depth: int) -> chess.engine.SimpleEngine: + """Create a Stockfish engine with given skill level.""" + engine = chess.engine.SimpleEngine.popen_uci(sf_path) + engine.configure({"Skill Level": skill}) + return engine + + +# --- Core benchmark functions --- + + +def play_model_vs_stockfish_game( + model: ChessTransformerV2, + device: torch.device, + engine: chess.engine.SimpleEngine, + sf_depth: int, + model_is_white: bool, + game_seed: int, +) -> str: + """Play one game, return '1-0', '0-1', or '1/2-1/2'.""" + rng = random.Random(game_seed) + opening = sample_opening(rng) + + board = chess.Board() + for uci in opening[3]: + move = chess.Move.from_uci(uci) + if move not in board.legal_moves: + break + board.push(move) + + with torch.no_grad(): + while not board.is_game_over(claim_draw=True) and board.ply() < 200: + is_white_turn = board.turn == chess.WHITE + + if is_white_turn == model_is_white: + board_t, feat_t = preprocess_board(board, device) + policy, promo, _wdl, _ply = model(board_t, feat_t) + move = greedy_move_v2(board, policy[0], promo[0]) + else: + result = engine.play(board, chess.engine.Limit(depth=sf_depth)) + move = result.move + + board.push(move) + + outcome = board.outcome(claim_draw=True) + if outcome is None or outcome.winner is None: + return "1/2-1/2" + return "1-0" if outcome.winner == chess.WHITE else "0-1" + + +def play_model_vs_model_game( + model_a: ChessTransformerV2, + model_b: ChessTransformerV2, + device: torch.device, + a_is_white: bool, + game_seed: int, +) -> str: + """Play one game between two models, return result string.""" + rng = random.Random(game_seed) + opening = sample_opening(rng) + + board = chess.Board() + for uci in opening[3]: + move = chess.Move.from_uci(uci) + if move not in board.legal_moves: + break + board.push(move) + + with torch.no_grad(): + while not board.is_game_over(claim_draw=True) and board.ply() < 200: + is_white_turn = board.turn == chess.WHITE + active_model = model_a if is_white_turn == a_is_white else model_b + + board_t, feat_t = preprocess_board(board, device) + policy, promo, _wdl, _ply = active_model(board_t, feat_t) + move = greedy_move_v2(board, policy[0], promo[0]) + board.push(move) + + outcome = board.outcome(claim_draw=True) + if outcome is None or outcome.winner is None: + return "1/2-1/2" + return "1-0" if outcome.winner == chess.WHITE else "0-1" + + +def benchmark_vs_stockfish( + model: ChessTransformerV2, + device: torch.device, + sf_path: str, + skill: int, + num_games: int, + sf_depth: int = 10, +) -> MatchResult: + """Run a match: model vs Stockfish at given skill level.""" + engine = _make_engine(sf_path, skill, sf_depth) + + wins, draws, losses = 0, 0, 0 + + try: + for i in range(num_games): + model_is_white = i % 2 == 0 + result = play_model_vs_stockfish_game( + model, device, engine, sf_depth, model_is_white, game_seed=i + ) + + if result == "1/2-1/2": + draws += 1 + elif (result == "1-0") == model_is_white: + wins += 1 + else: + losses += 1 + + if (i + 1) % 10 == 0 or i + 1 == num_games: + total = i + 1 + score = (wins + 0.5 * draws) / total + print(f" [{total}/{num_games}] +{wins} ={draws} -{losses} ({score:.1%})") + finally: + engine.quit() + + return MatchResult( + wins=wins, + draws=draws, + losses=losses, + label_a="Model", + label_b=f"Stockfish (skill {skill}, depth {sf_depth})", + ) + + +def benchmark_vs_model( + model_a: ChessTransformerV2, + model_b: ChessTransformerV2, + device: torch.device, + num_games: int, + label_a: str = "Model A", + label_b: str = "Model B", +) -> MatchResult: + """Run a match between two models.""" + model_a.eval() + model_b.eval() + wins, draws, losses = 0, 0, 0 + + for i in range(num_games): + a_is_white = i % 2 == 0 + result = play_model_vs_model_game( + model_a, model_b, device, a_is_white, game_seed=i + ) + + if result == "1/2-1/2": + draws += 1 + elif (result == "1-0") == a_is_white: + wins += 1 + else: + losses += 1 + + if (i + 1) % 10 == 0 or i + 1 == num_games: + total = i + 1 + score = (wins + 0.5 * draws) / total + print(f" [{total}/{num_games}] +{wins} ={draws} -{losses} ({score:.1%})") + + return MatchResult( + wins=wins, draws=draws, losses=losses, + label_a=label_a, label_b=label_b, + ) + + +def full_benchmark( + model: ChessTransformerV2, + device: torch.device, + sf_path: str, + num_games: int, + skills: tuple[int, ...] = (1, 3, 5, 7, 10), + sf_depth: int = 10, +) -> list[MatchResult]: + """Run model against multiple Stockfish skill levels.""" + results = [] + for skill in skills: + print(f"\n--- Stockfish Skill {skill} ({num_games} games) ---") + result = benchmark_vs_stockfish( + model, device, sf_path, skill, num_games, sf_depth + ) + results.append(result) + print(f" {result.summary()}") + return results + + +# --- CLI --- + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Benchmark ChessTransformerV2 models" + ) + subparsers = parser.add_subparsers(dest="command", required=True) + + # vs-stockfish + p_sf = subparsers.add_parser("vs-stockfish", help="Model vs Stockfish") + p_sf.add_argument("model", help="Path to V2 model checkpoint") + p_sf.add_argument("--skill", type=int, default=5, help="Stockfish skill level (0-20)") + p_sf.add_argument("--depth", type=int, default=10, help="Stockfish search depth") + p_sf.add_argument("--games", type=int, default=30, help="Number of games") + p_sf.add_argument("--sf-path", default=None, help="Path to Stockfish binary") + p_sf.add_argument("--device", default="auto", choices=["auto", "cuda", "mps", "cpu"]) + + # vs-model + p_mm = subparsers.add_parser("vs-model", help="Model vs Model") + p_mm.add_argument("model_a", help="Path to first model") + p_mm.add_argument("model_b", help="Path to second model") + p_mm.add_argument("--games", type=int, default=50, help="Number of games") + p_mm.add_argument("--device", default="auto", choices=["auto", "cuda", "mps", "cpu"]) + + # full + p_full = subparsers.add_parser("full", help="Full benchmark (multiple SF levels)") + p_full.add_argument("model", help="Path to V2 model checkpoint") + p_full.add_argument("--games", type=int, default=20, help="Games per skill level") + p_full.add_argument("--depth", type=int, default=10, help="Stockfish search depth") + p_full.add_argument("--skills", default="1,3,5,7,10", + help="Comma-separated skill levels (default: 1,3,5,7,10)") + p_full.add_argument("--sf-path", default=None, help="Path to Stockfish binary") + p_full.add_argument("--device", default="auto", choices=["auto", "cuda", "mps", "cpu"]) + + args = parser.parse_args() + device = detect_device(args.device) + + if args.command == "vs-stockfish": + sf_path = args.sf_path or _find_stockfish() + model, version, _ = load_model(args.model, device) + model.eval() + print(f"Model: {args.model} ({version})") + print(f"Stockfish: skill {args.skill}, depth {args.depth}") + print(f"Device: {device}") + print(f"Games: {args.games}\n") + + start = time.time() + result = benchmark_vs_stockfish( + model, device, sf_path, args.skill, args.games, args.depth + ) + elapsed = time.time() - start + print(f"\n{result.summary()}") + print(f"Time: {elapsed:.0f}s ({elapsed / args.games:.1f}s/game)") + + elif args.command == "vs-model": + model_a, ver_a, _ = load_model(args.model_a, device) + model_b, ver_b, _ = load_model(args.model_b, device) + model_a.eval() + model_b.eval() + label_a = args.model_a.split("/")[-1] + label_b = args.model_b.split("/")[-1] + print(f"Model A: {args.model_a} ({ver_a})") + print(f"Model B: {args.model_b} ({ver_b})") + print(f"Device: {device}") + print(f"Games: {args.games}\n") + + start = time.time() + result = benchmark_vs_model( + model_a, model_b, device, args.games, label_a, label_b + ) + elapsed = time.time() - start + print(f"\n{result.summary()}") + print(f"Time: {elapsed:.0f}s ({elapsed / args.games:.1f}s/game)") + + elif args.command == "full": + sf_path = args.sf_path or _find_stockfish() + model, version, _ = load_model(args.model, device) + model.eval() + skills = tuple(int(s) for s in args.skills.split(",")) + print(f"Model: {args.model} ({version})") + print(f"Stockfish levels: {skills}") + print(f"Device: {device}") + print(f"Games per level: {args.games}") + + start = time.time() + results = full_benchmark( + model, device, sf_path, args.games, skills, args.depth + ) + elapsed = time.time() - start + + print(f"\n{'='*60}") + print(" BENCHMARK RESULTS") + print(f"{'='*60}") + for r in results: + print(f" {r.summary()}") + print(f"{'='*60}") + print(f"Total time: {elapsed:.0f}s") + + +if __name__ == "__main__": + main() diff --git a/chess_loader.py b/chess_loader.py index cea8847..b6efc68 100644 --- a/chess_loader.py +++ b/chess_loader.py @@ -1,6 +1,23 @@ import torch from torch.utils.data import Dataset, DataLoader + +# --- V2 constants --- + +PIECE_TO_INDEX = { + '.': 0, 'P': 1, 'N': 2, 'B': 3, 'R': 4, 'Q': 5, 'K': 6, + 'p': 7, 'n': 8, 'b': 9, 'r': 10, 'q': 11, 'k': 12, +} + +# Piece values for material balance (uppercase = current player, lowercase = opponent) +PIECE_VALUES = { + 'P': 1, 'N': 3, 'B': 3, 'R': 5, 'Q': 9, + 'p': -1, 'n': -3, 'b': -3, 'r': -5, 'q': -9, +} +MAX_MATERIAL = 39.0 # Q + 2R + 2B + 2N + 8P + +PROMO_PIECES = {'q': 0, 'r': 1, 'b': 2, 'n': 3} + def square_num(sq: str) -> int: """ Converts chess square notation to a numerical index. @@ -26,9 +43,7 @@ def parse_pos_lists(list_file, num_pos=None): continue board, new_move = line.strip().split() - piece_to_index = {'.': 1, 'P': 2, 'N': 3, 'B': 4, 'R': 5, 'Q': 6, 'K': 7, - 'p': 8, 'n': 9, 'b': 10, 'r': 11, 'q': 12, 'k': 13} - board = [piece_to_index[p] for p in board] # Convert pieces to integers + board = [PIECE_TO_INDEX[p] for p in board] new_move = new_move[:2], new_move[2:] # Split move into start and end squares new_move = square_num(new_move[0]), square_num(new_move[1]) # Convert squares to indices @@ -64,11 +79,143 @@ def get_dataloader(pos_file, batch_size=32, num_workers=0, num_pos=None): test_len = min(5000, int(len(dataset) * 0.1)) dataset, testset = torch.utils.data.random_split(dataset, [len(dataset) - test_len, test_len]) + # Changed: pin_memory=True for faster CPU→GPU transfer (DMA direct copy) dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=num_workers) testloader = DataLoader(testset, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=num_workers) return dataloader, testloader + +# --- V2: Extended data pipeline with features and WDL --- + + +def compute_features(board_str: str) -> list[float]: + """Compute auxiliary features from 64-char board string. + + Returns 14 floats: material_balance(1) + check(1) + castling(4) + en_passant(8). + Only material_balance is computable from board string alone; + remaining features require FEN-level info and default to zero. + """ + material = sum(PIECE_VALUES.get(c, 0) for c in board_str) + return [material / MAX_MATERIAL] + [0.0] * 13 + + +def result_to_wdl(result: float) -> tuple[float, float, float]: + """Convert game result to WDL one-hot (from current player's perspective). + + 1.0 → win, 0.5 → draw, 0.0 → loss. + """ + if result > 0.75: + return (1.0, 0.0, 0.0) + if result > 0.25: + return (0.0, 1.0, 0.0) + return (0.0, 0.0, 1.0) + + +def parse_move(move_str: str) -> tuple[int, int, int]: + """Parse UCI move to (from_sq, to_sq, promo_idx). + + promo_idx: 0-3 for Q/R/B/N promotion, -1 if no promotion. + """ + from_sq = square_num(move_str[:2]) + to_sq = square_num(move_str[2:4]) + promo_idx = PROMO_PIECES.get(move_str[4:5], -1) + return from_sq, to_sq, promo_idx + + +def parse_pos_lists_v2( + list_file: str, num_pos: int | None = None +) -> tuple[list, list, list, list]: + """Parse position file with optional result column. + + Supports both formats: + <64-char board> (V1 compat, uniform WDL) + <64-char board> (V2 with game outcome) + + Returns: (boards, moves, features, wdl_targets) + """ + if isinstance(num_pos, float): + num_pos = int(num_pos) + if not isinstance(num_pos, int): + num_pos = int(1e9) + + boards: list[list[int]] = [] + moves: list[tuple[int, int, int]] = [] + features: list[list[float]] = [] + wdl_targets: list[tuple[float, float, float]] = [] + + with open(list_file, 'r') as file: + for i, line in enumerate(file): + if i >= num_pos: + break + line = line.strip() + if not line: + continue + + parts = line.split() + board_str = parts[0] + move_str = parts[1] + result = float(parts[2]) if len(parts) > 2 else None + + boards.append([PIECE_TO_INDEX[p] for p in board_str]) + moves.append(parse_move(move_str)) + features.append(compute_features(board_str)) + wdl_targets.append( + result_to_wdl(result) if result is not None + else (1 / 3, 1 / 3, 1 / 3) + ) + + return boards, moves, features, wdl_targets + + +class ChessDatasetV2(Dataset): + """Dataset for V2 model: board + features + from/to targets + WDL.""" + + def __init__(self, boards, moves, features, wdl_targets): + self.boards = torch.tensor(boards, dtype=torch.long) + self.from_sq = torch.tensor([m[0] for m in moves], dtype=torch.long) + self.to_sq = torch.tensor([m[1] for m in moves], dtype=torch.long) + self.promo = torch.tensor([m[2] for m in moves], dtype=torch.long) + self.features = torch.tensor(features, dtype=torch.float32) + self.wdl = torch.tensor(wdl_targets, dtype=torch.float32) + + def __len__(self): + return len(self.boards) + + def __getitem__(self, idx): + return ( + self.boards[idx], + self.features[idx], + self.from_sq[idx], + self.to_sq[idx], + self.promo[idx], + self.wdl[idx], + ) + + +def get_dataloader_v2(pos_file, batch_size=32, num_workers=0, num_pos=None): + """Creates train/test dataloaders for ChessTransformerV2.""" + boards, moves, features, wdl_targets = parse_pos_lists_v2( + pos_file, num_pos=num_pos + ) + dataset = ChessDatasetV2(boards, moves, features, wdl_targets) + + test_len = min(5000, int(len(dataset) * 0.1)) + dataset, testset = torch.utils.data.random_split( + dataset, [len(dataset) - test_len, test_len] + ) + + dataloader = DataLoader( + dataset, batch_size=batch_size, shuffle=True, + pin_memory=True, num_workers=num_workers, + ) + testloader = DataLoader( + testset, batch_size=batch_size, shuffle=True, + pin_memory=True, num_workers=num_workers, + ) + return dataloader, testloader + + if __name__ == '__main__': # Example usage of the get_dataloader function dataloader, testloader = get_dataloader('path_to_your_pgn_file.txt') diff --git a/chessformer.py b/chessformer.py index 9cc5123..556832d 100644 --- a/chessformer.py +++ b/chessformer.py @@ -3,21 +3,25 @@ from torch import nn, Tensor from torch.nn import TransformerEncoder, TransformerEncoderLayer +from attention import ShawRelativePositionBias, SmolgenBias, ChessTransformerBlock + class PositionalEncoding(nn.Module): def __init__(self, d_model: int, dropout: float = 0.1, max_len: int = 5000): super().__init__() self.dropout = nn.Dropout(p=dropout) # Creating positional encodings + # Changed: shape [1, max_len, d_model] for batch_first format (was [max_len, 1, d_model]) position = torch.arange(max_len).unsqueeze(1) div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)) - pe = torch.zeros(max_len, 1, d_model) - pe[:, 0, 0::2] = torch.sin(position * div_term) - pe[:, 0, 1::2] = torch.cos(position * div_term) + pe = torch.zeros(1, max_len, d_model) + pe[0, :, 0::2] = torch.sin(position * div_term) + pe[0, :, 1::2] = torch.cos(position * div_term) self.register_buffer('pe', pe) def forward(self, x: Tensor) -> Tensor: - x = x + self.pe[:x.size(0)] + # Changed: index dim 1 instead of dim 0 for batch_first format + x = x + self.pe[:, :x.size(1)] return self.dropout(x) class ChessTransformer(nn.Module): @@ -29,7 +33,8 @@ def __init__(self, inp_dict: int, out_dict: int, d_model: int, nhead: int, d_hid self.pos_encoder = PositionalEncoding(d_model, dropout) # Transformer Encoder Layers - encoder_layers = TransformerEncoderLayer(d_model, nhead, d_hid, dropout) + # Changed: batch_first=True enables Flash Attention and removes need for permute() + encoder_layers = TransformerEncoderLayer(d_model, nhead, d_hid, dropout, batch_first=True) self.transformer_encoder = TransformerEncoder(encoder_layers, nlayers) # Embedding layers for input @@ -41,6 +46,12 @@ def __init__(self, inp_dict: int, out_dict: int, d_model: int, nhead: int, d_hid # Output linear layer self.linear_output = nn.Linear(d_model, out_dict) + # Changed: pre-compute x/y coordinate tensors (shape [1, 64]) once instead of every forward pass + # x_coords: [0,1,2,3,4,5,6,7, 0,1,2,3,4,5,6,7, ...] (column index for each square) + # y_coords: [0,0,0,0,0,0,0,0, 1,1,1,1,1,1,1,1, ...] (row index for each square) + self.register_buffer('x_coords', torch.arange(8).unsqueeze(0).repeat(1, 8)) + self.register_buffer('y_coords', torch.arange(8).repeat_interleave(8).unsqueeze(0)) + # Initialization of weights self.init_weights() @@ -55,21 +66,295 @@ def init_weights(self) -> None: def forward(self, board: Tensor, src_mask: Tensor = None) -> Tensor: board_emb = self.embedding(board) - # Generating input for positional embeddings - batch_size = board.shape[0] - x_inp = torch.arange(8).unsqueeze(0).repeat(batch_size, 8).to(board.device) - y_inp = torch.arange(8).repeat_interleave(8).unsqueeze(0).repeat(batch_size, 1).to(board.device) - - x_emb = self.x_embedding(x_inp) - y_emb = self.y_embedding(y_inp) + # Changed: use pre-computed coordinate buffers, expand() creates a view without copying memory + x_emb = self.x_embedding(self.x_coords.expand(board.shape[0], -1)) + y_emb = self.y_embedding(self.y_coords.expand(board.shape[0], -1)) # Combining embeddings combined_emb = board_emb + x_emb + y_emb # Scaling and passing through transformer combined_emb = combined_emb * math.sqrt(self.d_model) - output = self.transformer_encoder(combined_emb.permute(1, 0, 2)).permute(1, 0, 2) + # Changed: removed .permute() calls - no longer needed with batch_first=True + output = self.transformer_encoder(combined_emb) # Applying linear layer to the output output = self.linear_output(output) return output + + +# --- V2: Enhanced architecture with Shaw RPE, Smolgen, WDLP --- + +NUM_AUX_FEATURES = 14 # material(1) + check(1) + castling(4) + en_passant(8) + + +class SourceDestPolicyHead(nn.Module): + """Bilinear source-destination policy head. + + Computes attention between source-square queries and destination-square + keys, producing a [B, 64, 64] logit matrix where entry (i, j) scores + moving from square i to square j. + """ + + def __init__(self, d_model: int, d_policy: int = 128) -> None: + super().__init__() + self.from_proj = nn.Linear(d_model, d_policy) + self.to_proj = nn.Linear(d_model, d_policy) + self.scale = d_policy ** -0.5 + self.promo_head = nn.Linear(d_model, 4) # Q, R, B, N + + def forward(self, x: Tensor) -> tuple[Tensor, Tensor]: + """ + Args: + x: [B, 64, d_model] + Returns: + policy_logits: [B, 64, 64] — source-destination move scores + promo_logits: [B, 64, 4] — promotion piece scores per source square + """ + q = self.from_proj(x) + k = self.to_proj(x) + policy = torch.bmm(q, k.transpose(1, 2)) * self.scale + promo = self.promo_head(x) + return policy, promo + + +class WDLPValueHead(nn.Module): + """Win/Draw/Loss + Ply prediction value head. + + Mean-pools 64 latent tokens, then predicts: + - WDL: 3-class probability distribution (softmax) + - Ply: expected game length (non-negative via softplus, auxiliary signal) + """ + + def __init__(self, d_model: int) -> None: + super().__init__() + self.norm = nn.LayerNorm(d_model) + self.wdl_linear = nn.Linear(d_model, 3) + self.ply_linear = nn.Linear(d_model, 1) + + def forward(self, x: Tensor) -> tuple[Tensor, Tensor]: + """ + Args: + x: [B, 64, d_model] + Returns: + wdl: [B, 3] — win/draw/loss probabilities (sum to 1) + ply: [B, 1] — predicted game length (non-negative) + """ + pooled = self.norm(x.mean(dim=1)) + wdl = torch.softmax(self.wdl_linear(pooled), dim=-1) + ply = torch.nn.functional.softplus(self.ply_linear(pooled)) + return wdl, ply + + +class ChessTransformerV2(nn.Module): + """Enhanced chess transformer with Shaw RPE, smolgen, and dual heads. + + Improvements over V1: + - Shaw relative position encoding (topology-aware, not sinusoidal) + - Smolgen dynamic attention biases (content-dependent) + - Pre-norm transformer blocks (more stable training) + - Source-destination policy head (structured 64x64 logit matrix) + - WDLP value head (win/draw/loss + ply prediction) + - Auxiliary input features (material, check, castling, en passant) + + Phase 2 extension (optional): + - Diffusion inference path: backbone latent → DiT denoising → fused policy + - Activated by calling attach_diffusion() then forward(use_diffusion=True) + """ + + def __init__( + self, + d_model: int = 512, + nhead: int = 8, + d_hid: int = 1024, + nlayers: int = 12, + d_policy: int = 128, + dropout: float = 0.1, + ) -> None: + super().__init__() + self.model_type = "Chess-former-v2" + self.d_model = d_model + + # Input embeddings (same as V1) + self.embedding = nn.Embedding(13, d_model) + self.x_embedding = nn.Embedding(8, d_model) + self.y_embedding = nn.Embedding(8, d_model) + + # Auxiliary feature fusion: project (d_model + 14) → d_model + self.feature_proj = nn.Linear(d_model + NUM_AUX_FEATURES, d_model) + + # Pre-computed coordinate buffers (same as V1) + self.register_buffer( + "x_coords", torch.arange(8).unsqueeze(0).repeat(1, 8) + ) + self.register_buffer( + "y_coords", torch.arange(8).repeat_interleave(8).unsqueeze(0) + ) + + # Attention biases + self.shaw_rpe = ShawRelativePositionBias(nhead) + self.smolgen = SmolgenBias(d_model, nhead) + + # Transformer backbone (pre-norm blocks) + self.layers = nn.ModuleList( + [ChessTransformerBlock(d_model, nhead, d_hid) for _ in range(nlayers)] + ) + self.final_norm = nn.LayerNorm(d_model) + + # Dual output heads + self.policy_head = SourceDestPolicyHead(d_model, d_policy) + self.value_head = WDLPValueHead(d_model) + + self.dropout = nn.Dropout(dropout) + self._init_weights() + + def _init_weights(self) -> None: + initrange = 0.1 + self.embedding.weight.data.uniform_(-initrange, initrange) + self.x_embedding.weight.data.uniform_(-initrange, initrange) + self.y_embedding.weight.data.uniform_(-initrange, initrange) + + def encode( + self, + board: Tensor, + features: Tensor | None = None, + ) -> Tensor: + """Run backbone only, return latent representation. + + Args: + board: [B, 64] int tensor — piece indices (0-12) + features: [B, 14] float tensor — auxiliary features, or None + + Returns: + latent: [B, 64, d_model] — backbone output before heads + """ + B = board.shape[0] + + # Embeddings (same as V1) + board_emb = self.embedding(board) + x_emb = self.x_embedding(self.x_coords.expand(B, -1)) + y_emb = self.y_embedding(self.y_coords.expand(B, -1)) + combined = board_emb + x_emb + y_emb + + # Fuse auxiliary features if provided + if features is not None: + feat_expanded = features.unsqueeze(1).expand(-1, 64, -1) + combined = self.feature_proj( + torch.cat([combined, feat_expanded], dim=-1) + ) + + combined = combined * math.sqrt(self.d_model) + combined = self.dropout(combined) + + # Compute attention biases (once, shared across all layers) + shaw_bias = self.shaw_rpe() # [nhead, 64, 64] + smolgen_bias = self.smolgen(combined) # [B, nhead, 64, 64] + attn_bias = smolgen_bias + shaw_bias.unsqueeze(0) + + # Transformer backbone + x = combined + for layer in self.layers: + x = layer(x, attn_bias=attn_bias) + return self.final_norm(x) + + def forward( + self, + board: Tensor, + features: Tensor | None = None, + use_diffusion: bool = False, + ) -> tuple[Tensor, Tensor, Tensor, Tensor]: + """Full forward pass: backbone → (optional diffusion) → heads. + + Args: + board: [B, 64] int tensor — piece indices (0-12) + features: [B, 14] float tensor — auxiliary features, or None + use_diffusion: if True and diffusion is attached, run denoising + + Returns: + policy_logits: [B, 64, 64] — source-destination move scores + promo_logits: [B, 64, 4] — promotion piece scores + wdl: [B, 3] — win/draw/loss probabilities + ply: [B, 1] — predicted game length + """ + latent = self.encode(board, features) + + # Phase 2: diffusion-augmented inference + if use_diffusion and hasattr(self, "diffusion") and self.diffusion is not None: + latent = self._diffusion_augment(latent) + + # Dual heads + policy_logits, promo_logits = self.policy_head(latent) + wdl, ply = self.value_head(latent) + + return policy_logits, promo_logits, wdl, ply + + def attach_diffusion( + self, + diffusion: nn.Module, + noise_schedule: object, + d_dit: int, + ) -> None: + """Attach Phase 2 diffusion components for augmented inference. + + Creates projection layers (d_model ↔ d_dit) and stores the + diffusion model + noise schedule. Call forward(use_diffusion=True) + to activate the denoising path. + + Args: + diffusion: ChessDiT model (predicts noise) + noise_schedule: CosineNoiseSchedule (provides alpha_bar, etc.) + d_dit: diffusion latent dimension + """ + self.diffusion = diffusion + self.noise_schedule = noise_schedule + self.d_dit = d_dit + # Project backbone latent → d_dit and back + self.latent_to_dit = nn.Linear(self.d_model, d_dit) + self.dit_to_latent = nn.Linear(d_dit, self.d_model) + # Zero-init the back-projection so diffusion starts as no-op + nn.init.zeros_(self.dit_to_latent.weight) + nn.init.zeros_(self.dit_to_latent.bias) + + def _diffusion_augment(self, latent: Tensor) -> Tensor: + """Denoise imagined future state and fuse with current latent. + + Runs the full reverse diffusion process: starts from pure noise, + denoises T steps conditioned on the current backbone latent, + then adds the result as a residual. + + Args: + latent: [B, 64, d_model] — backbone output + + Returns: + augmented: [B, 64, d_model] — latent + diffusion context + """ + ns = self.noise_schedule + B = latent.shape[0] + device = latent.device + + # Start from pure noise in d_dit space + x_t = torch.randn(B, 64, self.d_dit, device=device) + + # Reverse diffusion (T → 0) + for t_val in reversed(range(ns.T)): + t_tensor = torch.full((B,), t_val, device=device, dtype=torch.long) + predicted_noise = self.diffusion(x_t, t_tensor, latent) + + # DDPM update step + alpha = ns.alpha[t_val] + beta = ns.beta[t_val] + + # Mean: (1/√α_t) * (x_t - (β_t/√(1-ᾱ_t)) * ε_θ) + coef = beta / ns.sqrt_one_minus_alpha_bar[t_val] + mean = (x_t - coef * predicted_noise) / alpha.sqrt() + + if t_val > 0: + # Add noise (not at final step) + sigma = ns.posterior_variance[t_val].sqrt() + x_t = mean + sigma * torch.randn_like(x_t) + else: + x_t = mean + + # Project back to d_model and fuse as residual + diff_context = self.dit_to_latent(x_t) + return latent + diff_context diff --git a/data_generation.md b/data_generation.md new file mode 100644 index 0000000..90321c0 --- /dev/null +++ b/data_generation.md @@ -0,0 +1,383 @@ +# Chess Training Data Generation — Theoretical Analysis & Implementation Plan + +## Status + +| Phase | Description | Status | +| --- | --- | --- | +| 1a | `openings.py` — curated opening FENs | ✅ Complete | +| 1b | `generate_selfplay_data.py` — sequential generation | ✅ Complete | +| 2 | Parallel generation via `multiprocessing.Pool` | ✅ Complete | +| 3 | GPU-accelerated encoding / lc0 integration | Planned | + +--- + +## Executive Summary + +This document outlines a strategy for generating high-quality, diverse chess training data using Stockfish self-play. The core challenge is that **two deterministic engines playing from the same position always produce the same game** — so we need principled randomness injection. The data is saved in PGN format and later converted to the existing training format. + +--- + +## 1. Theoretical Motivation + +### Why Self-Play Beats Lichess Data + +| Property | Lichess Data | Stockfish Self-Play | +| --- | --- | --- | +| Move quality | Mixed (humans blunder) | Controllable | +| Position diversity | Very high | Needs engineering | +| ELO distribution | Wide | Fully controlled | +| Annotation | None (raw moves) | Can embed evaluations | +| Data volume | Limited by database | Unlimited | +| Perspective | Human style | Engine style | + +Lichess data at 2000+ ELO still contains significant human error patterns, opening prep bias (players repeating memorized lines), and time-pressure blunders. Stockfish self-play at calibrated strength gives **precise, reproducible quality** and **no noise from time pressure or psychology**. + +The key insight from AlphaZero (DeepMind, 2017): a model trained exclusively on self-play data dramatically outperformed one trained on human games, even when the human dataset was vastly larger. + +### The Determinism Problem + +Stockfish is a deterministic engine: given the same position and the same parameters, it always plays the same move. Two games started from the initial position with identical settings will be **byte-for-byte identical**. This makes naive self-play useless — you would generate millions of copies of the same game. + +Solution: inject randomness at multiple levels to ensure diverse, non-repeating games. + +--- + +## 2. Randomness Injection Strategies + +### 2.1 Opening Diversification (Most Important) + +The opening phase determines the entire character of the game — if all games start the same, they all diverge at the same point and may converge to similar middlegame structures. + +#### Strategy A — Random Legal Moves for First N Plies + +Apply N (e.g., 4–8) uniformly random legal moves before handing control to Stockfish. This scatters games across a huge variety of positions quickly. + +- Pros: Extremely simple, maximum variety +- Cons: May produce strategically nonsensical positions that never appear in real play + +#### Strategy B — Weighted Random from Stockfish MultiPV (Recommended) + +Use Stockfish's `MultiPV` mode to compute the top-K moves, then sample one move proportionally to their evaluation scores (temperature sampling). This produces varied but *plausible* positions. + +```text +scores = [eval_1, eval_2, ..., eval_k] # centipawns +probs = softmax(scores / T) # temperature T controls spread +move = random.choices(top_k_moves, weights=probs) +``` + +With temperature T: + +- T → 0: Always play best move (deterministic) +- T = 50 cp: Strong preference for top moves, some variety +- T = 100 cp: Moderate randomness, all top-3 moves are plausible +- T → ∞: Uniform random (chaotic) + +For opening phase (moves 1–15): use T = 80–120 cp. +For middlegame (moves 15+): hand control to Stockfish deterministically. + +#### Strategy C — Curated Opening Book + +Maintain a list of starting FENs or move sequences (e.g., 500 common ECO openings) and pick one randomly per game. + +- Pros: Realistic positions, covers all major opening systems +- Cons: Requires curation; still deterministic after the book ends + +#### Strategy D — Polyglot Opening Book + +Use an existing Polyglot `.bin` opening book (e.g., `gm2001.bin`, `komodo.bin`). Follow the book for opening moves, weighted by book move frequency, then switch to Stockfish. + +#### Recommended Approach: Combine B + C + +- Start from a random opening FEN (from curated list of ~300 openings) +- Apply temperature sampling for moves 1–12 (relative to that opening) +- Switch to pure Stockfish for the rest of the game + +--- + +### 2.2 Engine Strength Variation + +Playing only at maximum strength (depth 20+) produces perfect play that's less instructive for the model — it won't see the kinds of positions that arise at 2000 ELO. Mix strength levels: + +#### Stockfish Skill Level (0–20) + +UCI parameter: `setoption name Skill Level value X` + +Approximate ELO mapping: + +| Skill Level | Approx ELO | +| --- | --- | +| 0 | ~1350 | +| 5 | ~1600 | +| 10 | ~1800 | +| 15 | ~2100 | +| 18 | ~2600 | +| 20 | ~3500+ | + +Internally, Skill Level works by: (a) limiting the depth, and (b) occasionally playing a random move weighted by evaluation. + +#### Direct ELO Limiting (More Precise) + +```text +setoption name UCI_LimitStrength value true +setoption name UCI_Elo value 2000 +``` + +UCI_Elo range: 1320–3190. This is the most principled way to target a specific strength. + +#### Depth Variation + +`go depth X` — search to fixed depth instead of time limit. + +- Depth 5–8: Fast, ~1800 ELO equivalent +- Depth 10–12: Medium, ~2200 ELO +- Depth 15: Strong, ~2600 ELO +- Depth 20+: Near-maximum, 3000+ ELO + +#### Recommended Strength Distribution for Training Data + +| Category | % of games | Config | +| --- | --- | --- | +| Elite | 20% | UCI_Elo=3000, depth=18 | +| Strong | 35% | UCI_Elo=2400, depth=14 | +| Club | 30% | UCI_Elo=2000, depth=10 | +| Intermediate | 15% | UCI_Elo=1600, depth=7 | + +This distribution mirrors real Lichess ratings and ensures the model learns patterns at multiple strength levels. + +--- + +### 2.3 Move Time and Node Variation + +Instead of fixed depth, use `go movetime X` (milliseconds) or `go nodes X`. This produces natural time-pressure variation: + +- Fast games (100ms/move): More tactical errors, similar to blitz +- Slow games (2000ms/move): Very accurate, near-optimal +- Node limits: More CPU-consistent across machines + +--- + +### 2.4 Asymmetric Games + +Play White and Black at different strengths: + +- Strong White vs. Weak Black: Model learns to press advantages, convert wins +- Weak White vs. Strong Black: Model learns defense, resource saving +- Equal strength: Most common, balanced games + +This produces a dataset with varied game outcomes (wins, draws, losses) rather than all games being drawn (which pure max-strength Stockfish games often are). + +--- + +## 3. Data Quality for Transformer Training + +### What Transformers Learn Best From + +1. **Move prediction**: Learn the mapping `position → best_move`. Need high-quality moves. +2. **Position understanding**: The board representation must cover all common position types. +3. **Diversity**: The model must see varied pawn structures, piece configurations, and endgame types. + +### Key Data Quality Metrics + +- **Position uniqueness**: What fraction of positions in the dataset are unique? Target >90%. +- **Phase coverage**: Ratio of opening/middlegame/endgame positions. Lichess data is opening-heavy; self-play can be tuned. +- **Centipawn variance**: A dataset of only equal positions (±50 cp) doesn't teach the model to handle decisive advantages. Mix in won/lost positions. +- **Move diversity per position**: Ensure the dataset doesn't repeat the same move from similar positions too often. + +### Evaluation Annotations + +Unlike raw Lichess PGN, self-play allows embedding Stockfish evaluations with every move. This unlocks potential future training signals: + +```pgn +1. e4 {+0.35/14} e5 {+0.15/14} 2. Nf3 {+0.42/14} ... +``` + +Even if the current model ignores evaluations, recording them now means the data can be reused for future value-head training (as in AlphaZero). + +--- + +## 4. PGN Format Specification + +### Standard PGN Headers + +```pgn +[Event "SF-SelfPlay"] +[Site "local"] +[Date "2026.02.18"] +[Round "1"] +[White "Stockfish_2000"] +[Black "Stockfish_1600"] +[Result "1-0"] +[WhiteElo "2000"] +[BlackElo "1600"] +[Opening "Sicilian Defense"] +[ECO "B20"] +[TimeControl "-"] +[Termination "Normal"] +[SFDepthWhite "10"] +[SFDepthBlack "7"] +``` + +Custom headers `SFDepthWhite`, `SFDepthBlack` track engine settings for reproducibility. + +### Optional Evaluation Annotations + +```pgn +1. e4 { [%eval 0.35] } 1... c5 { [%eval 0.28] } 2. Nf3 { [%eval 0.41] } +``` + +This is the standard Lichess PGN annotation format (compatible with `pgn_to_training_data.py`). + +--- + +## 5. Implementation + +### Phase 1a — Opening Book (`openings.py`) ✅ + +File: `openings.py` + +Contains ~250 curated ECO openings as `(eco_code, name, fen, moves_uci)` tuples covering all major opening systems: Open Games, Sicilian, French, Caro-Kann, Queen's Gambit, King's Indian, English, Nimzo-Indian, and more. + +```python +def load_openings() -> list[tuple]: + """Return list of (eco, name, fen, moves_uci) for all openings.""" + return OPENINGS +``` + +### Phase 1b — Sequential Generator (`generate_selfplay_data.py`) ✅ + +File: `generate_selfplay_data.py` + +Core components: + +- `GameConfig` dataclass — all parameters for one game +- `temperature_sample(multipv_info, temperature_cp)` — softmax sampling over top-K moves +- `generate_game(sf_path, config, game_id)` — produces one `chess.pgn.Game` +- `sample_game_config(openings)` — randomly samples a `GameConfig` +- `main()` — sequential CLI entry point + +```bash +python generate_selfplay_data.py \ + --sf-path ./stockfish/stockfish-ubuntu-x86-64-avx2 \ + --output selfplay_data/run_001.pgn \ + --num-games 500 \ + --seed 42 +``` + +### Phase 2 — Parallel Generator ✅ + +Same file `generate_selfplay_data.py`, add `--workers N` flag. + +Uses `multiprocessing.Pool` with `imap_unordered` — each worker owns an independent Stockfish process. Expected speedup: ~linear with core count. + +```bash +python generate_selfplay_data.py \ + --sf-path ./stockfish/stockfish-ubuntu-x86-64-avx2 \ + --output selfplay_data/parallel_run.pgn \ + --num-games 5000 \ + --workers 8 \ + --seed 42 +``` + +### Phase 3 — GPU Parallelisation (Planned) + +GPU acceleration for chess self-play is unconventional — Stockfish is CPU-only. The GPU is useful for: + +1. **Neural network inference in the loop**: Mix our model's moves with Stockfish for curriculum learning; GPU inference can be batched. +2. **Leela Chess Zero (lc0)**: GPU-accelerated neural engine as generator, produces more human-like patterns. +3. **CUDA-accelerated game encoding**: Convert PGN → training tensors on GPU rather than CPU. + +lc0 natively supports temperature sampling: + +```text +setoption name Temperature value 0.8 +setoption name TempDecayMoves value 30 +``` + +--- + +## 6. Integration with Existing Pipeline + +### PGN → Training Data + +`pgn_to_training_data.py` was extended to skip ELO/time-control filters for self-play games (detected via `Event` header starting with `"SF-SelfPlay"`). + +### Full Data Pipeline + +```text +generate_selfplay_data.py + --num-games 10000 + --output selfplay_data/raw.pgn + ↓ +pgn_to_training_data.py + selfplay_data/raw.pgn + full_datasets/selfplay_pos.txt + 0 (min_elo=0, no filter for self-play) + ↓ +train_model.py + --dataset full_datasets/selfplay_pos.txt + --model models/selfplay_v1.pth +``` + +--- + +## 7. Expected Data Volume & Quality + +### Games per Hour Estimates (CPU) + +| Config | Games/hr (1 core) | Games/hr (8 cores) | +| --- | --- | --- | +| Depth 7, fast openings | ~180 | ~1400 | +| Depth 10, mixed | ~80 | ~620 | +| Depth 14, annotated | ~25 | ~190 | +| Depth 18, elite | ~8 | ~60 | + +### Positions per Game (Average) + +Average game length: 40–80 moves → 80–160 positions per game (both sides). + +- 1,000 games ≈ 100,000 positions +- 10,000 games ≈ 1,000,000 positions (matches current Lichess dataset size) +- 100,000 games ≈ 10,000,000 positions + +### Data Quality Comparison + +| Metric | Lichess 2000+ | SF Self-Play (2000 ELO) | SF Self-Play (mixed) | +| --- | --- | --- | --- | +| Avg centipawn loss | ~50 | ~10 | ~25 | +| Position uniqueness | ~98% | ~95% | ~97% | +| Phase coverage | Opening-heavy | Balanced | Balanced | +| Move consistency | Human-like | Engine-like | Engine-like | + +--- + +## 8. Key Risks and Mitigations + +| Risk | Mitigation | +| --- | --- | +| Games too similar (low diversity) | Combine opening book + temperature sampling | +| Stockfish process leaks | Always use `engine.quit()` in try/finally blocks | +| PGN parse errors | Validate with `python-chess` before writing | +| Drawn games dominate at high ELO | Use asymmetric strengths; target 40-30-30 W/D/L split | +| CPU bottleneck | Use multiprocessing (Phase 2); prioritize fast depth settings | +| Training worse than Lichess data | Mix self-play + Lichess data; A/B test model quality | + +--- + +## 9. File Structure + +```text +chessformer/ +├── generate_selfplay_data.py # Main generator (Phase 1 + 2) +├── openings.py # Opening book (~250 ECO openings) +├── pgn_to_training_data.py # Extended to handle self-play PGN +├── selfplay_data/ # Output directory (gitignored) +│ ├── run_001.pgn +│ └── ... +└── full_datasets/ + ├── elo_2000_pos.txt # Existing Lichess data + └── selfplay_pos.txt # Self-play converted data +``` + +--- + +Document version: 2.0 — Feb 2026 diff --git a/diffusion_model.py b/diffusion_model.py new file mode 100644 index 0000000..b324d6e --- /dev/null +++ b/diffusion_model.py @@ -0,0 +1,207 @@ +"""diffusion_model.py — DiT (Diffusion Transformer) for chess latent space. + +Denoises "imagined" future board states in the latent space of the policy +backbone. Conditioned on the current position's latent representation and +the diffusion timestep. + +Architecture: + - Input: noisy latent x_t [B, 64, d_dit] + - Conditioning: backbone latent [B, 64, d_model] → mean-pooled → projected + - Timestep: Embedding(T, d_dit) → MLP + - DiT blocks with AdaLN-Zero (adaptive LayerNorm + zero-init gates) + - Output: predicted noise epsilon [B, 64, d_dit] + +AdaLN-Zero trick: gate parameters initialize to 0, so each block starts +as identity. This means the model begins training as a simple pass-through, +then gradually learns to denoise — much more stable than random init. + +Reference: Peebles & Xie "Scalable Diffusion Models with Transformers" (2023). +""" + +import torch +from torch import Tensor, nn +from torch.nn import functional as F + + +class TimestepEmbedding(nn.Module): + """Sinusoidal + MLP timestep embedding. + + Maps scalar timestep t → d_dit vector. Uses sinusoidal encoding + (same as positional encoding in original Transformer) followed by + a 2-layer MLP to project to the right dimension. + """ + + def __init__(self, T: int, d_dit: int) -> None: + super().__init__() + self.embed = nn.Embedding(T + 1, d_dit) + self.mlp = nn.Sequential( + nn.Linear(d_dit, d_dit * 4), + nn.Mish(), + nn.Linear(d_dit * 4, d_dit), + ) + + def forward(self, t: Tensor) -> Tensor: + """[B] int → [B, d_dit] float.""" + return self.mlp(self.embed(t)) + + +class DiTBlock(nn.Module): + """Transformer block with AdaLN-Zero conditioning. + + Standard transformer block (self-attention + FFN) but LayerNorm parameters + are generated from a conditioning vector c. Additionally, each sub-block + has a learned gate that starts at 0 (zero-init), making the block act as + identity at initialization. + + Conditioning c produces 6 vectors via a single linear projection: + shift1, scale1, gate1 (for attention) + shift2, scale2, gate2 (for FFN) + + Then: + h = LN(x) * (1 + scale1) + shift1 + h = self_attention(h) + x = x + gate1 * h + h = LN(x) * (1 + scale2) + shift2 + h = ffn(h) + x = x + gate2 * h + """ + + def __init__(self, d_dit: int, num_heads: int, d_hid: int) -> None: + super().__init__() + self.num_heads = num_heads + self.head_dim = d_dit // num_heads + + # AdaLN modulation: c → 6 * d_dit params + self.adaLN_modulation = nn.Sequential( + nn.Mish(), + nn.Linear(d_dit, 6 * d_dit), + ) + # Zero-init the final linear so gates start at 0 + nn.init.zeros_(self.adaLN_modulation[1].weight) + nn.init.zeros_(self.adaLN_modulation[1].bias) + + self.norm1 = nn.LayerNorm(d_dit, elementwise_affine=False) + self.qkv = nn.Linear(d_dit, 3 * d_dit) + self.out_proj = nn.Linear(d_dit, d_dit) + + self.norm2 = nn.LayerNorm(d_dit, elementwise_affine=False) + self.ffn = nn.Sequential( + nn.Linear(d_dit, d_hid), + nn.Mish(), + nn.Linear(d_hid, d_dit), + ) + + def forward(self, x: Tensor, c: Tensor) -> Tensor: + """ + Args: + x: [B, 64, d_dit] — noisy latent tokens + c: [B, d_dit] — conditioning vector (timestep + position) + """ + B, S, D = x.shape + + # Generate modulation params from conditioning + mod = self.adaLN_modulation(c).unsqueeze(1) # [B, 1, 6*d_dit] + shift1, scale1, gate1, shift2, scale2, gate2 = mod.chunk(6, dim=-1) + + # Modulated self-attention + h = self.norm1(x) * (1 + scale1) + shift1 + qkv = self.qkv(h).reshape(B, S, 3, self.num_heads, self.head_dim) + q, k, v = qkv.unbind(dim=2) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + attn = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5) + attn = F.softmax(attn, dim=-1) + h = (attn @ v).transpose(1, 2).reshape(B, S, D) + h = self.out_proj(h) + x = x + gate1 * h + + # Modulated FFN + h = self.norm2(x) * (1 + scale2) + shift2 + h = self.ffn(h) + x = x + gate2 * h + + return x + + +class ChessDiT(nn.Module): + """Diffusion Transformer for chess latent-space denoising. + + Takes a noisy latent state x_t and predicts the noise epsilon that was + added, conditioned on the current position's latent representation and + the diffusion timestep. + + Config: d_dit=256, nhead=4, nlayers=6, T=20 → ~10M params + """ + + def __init__( + self, + d_dit: int = 256, + d_model: int = 512, + nhead: int = 4, + d_hid: int = 512, + nlayers: int = 6, + T: int = 20, + ) -> None: + super().__init__() + self.d_dit = d_dit + + # Timestep embedding + self.time_embed = TimestepEmbedding(T, d_dit) + + # Condition projection: backbone latent → d_dit + # Mean-pools 64 tokens to a single vector, then projects + self.cond_proj = nn.Sequential( + nn.Linear(d_model, d_dit), + nn.Mish(), + nn.Linear(d_dit, d_dit), + ) + + # DiT blocks + self.layers = nn.ModuleList( + [DiTBlock(d_dit, nhead, d_hid) for _ in range(nlayers)] + ) + + # Final layer: LN + linear, zero-initialized for stable start + self.final_norm = nn.LayerNorm(d_dit, elementwise_affine=False) + self.final_adaLN = nn.Sequential(nn.Mish(), nn.Linear(d_dit, 2 * d_dit)) + nn.init.zeros_(self.final_adaLN[1].weight) + nn.init.zeros_(self.final_adaLN[1].bias) + self.final_proj = nn.Linear(d_dit, d_dit) + # Zero-init final projection so initial output is zero (pure residual) + nn.init.zeros_(self.final_proj.weight) + nn.init.zeros_(self.final_proj.bias) + + def forward( + self, + x_t: Tensor, + t: Tensor, + backbone_latent: Tensor, + ) -> Tensor: + """Predict noise epsilon from noisy latent. + + Args: + x_t: [B, 64, d_dit] — noisy latent state + t: [B] int — diffusion timestep + backbone_latent: [B, 64, d_model] — current position encoding from backbone + + Returns: + epsilon: [B, 64, d_dit] — predicted noise + """ + # Combine timestep + condition into a single vector + time_emb = self.time_embed(t) # [B, d_dit] + cond = self.cond_proj(backbone_latent.mean(dim=1)) # [B, d_dit] + c = time_emb + cond # [B, d_dit] + + # Pass through DiT blocks + x = x_t + for layer in self.layers: + x = layer(x, c) + + # Final layer with AdaLN + mod = self.final_adaLN(c).unsqueeze(1) # [B, 1, 2*d_dit] + shift, scale = mod.chunk(2, dim=-1) + x = self.final_norm(x) * (1 + scale) + shift + epsilon = self.final_proj(x) + + return epsilon diff --git a/fonts/NotoSansSymbols2-Regular.ttf b/fonts/NotoSansSymbols2-Regular.ttf new file mode 100644 index 0000000..45df830 Binary files /dev/null and b/fonts/NotoSansSymbols2-Regular.ttf differ diff --git a/generate_selfplay_data.py b/generate_selfplay_data.py new file mode 100644 index 0000000..cad4a1e --- /dev/null +++ b/generate_selfplay_data.py @@ -0,0 +1,396 @@ +""" +Stockfish self-play data generator. + +Generates diverse chess games between Stockfish instances at varying strength +levels and saves them in PGN format for use as training data. + +Randomness is injected via: + 1. Opening book — games start from ~250 curated ECO opening positions. + 2. Temperature sampling — first temp_cutoff_ply moves are sampled from + Stockfish MultiPV results via softmax, producing varied but plausible play. + 3. Asymmetric strength — White and Black are assigned different ELO targets + and search depths per game, producing a mix of outcomes. + +Usage: + # Sequential (Phase 1) + python generate_selfplay_data.py \\ + --sf-path ./stockfish/stockfish-ubuntu-x86-64-avx2 \\ + --output selfplay_data/run_001.pgn \\ + --num-games 500 --seed 42 + + # Parallel (Phase 2) + python generate_selfplay_data.py \\ + --sf-path ./stockfish/stockfish-ubuntu-x86-64-avx2 \\ + --output selfplay_data/run_parallel.pgn \\ + --num-games 5000 --workers 8 --seed 42 +""" + +from __future__ import annotations + +import argparse +import os +import random +import sys +import time +from dataclasses import dataclass +from datetime import date +from multiprocessing import Pool + +import chess +import chess.engine +import chess.pgn +import numpy as np + +from openings import load_openings, sample_opening + +# --------------------------------------------------------------------------- +# Strength tiers: (weight, white_elo, black_elo, white_depth, black_depth) +# --------------------------------------------------------------------------- +_STRENGTH_TIERS = [ + (0.15, 3000, 3000, 14, 14), # Elite vs Elite + (0.10, 2800, 2400, 13, 11), # Asymmetric strong (W > B) + (0.10, 2400, 2800, 11, 13), # Asymmetric strong (B > W) + (0.20, 2200, 2200, 11, 11), # Strong club + (0.20, 2000, 2000, 9, 9), # Club + (0.10, 2000, 1600, 9, 7), # Intermediate asymmetric + (0.10, 1600, 2000, 7, 9), # Intermediate asymmetric (flipped) + (0.05, 1600, 1600, 7, 7), # Intermediate +] +_TIER_WEIGHTS = [t[0] for t in _STRENGTH_TIERS] + +# UCI_Elo is clamped to Stockfish's supported range +_UCI_ELO_MIN = 1320 +_UCI_ELO_MAX = 3190 + + +@dataclass +class GameConfig: + """All parameters needed to generate one self-play game.""" + white_elo: int + black_elo: int + white_depth: int + black_depth: int + opening_eco: str + opening_name: str + opening_fen: str + opening_moves: list[str] + temperature: float # centipawns; controls opening move diversity + temp_cutoff_ply: int # stop temperature sampling after this many plies + annotate_evals: bool + game_id: int + seed: int + + +def _clamp_elo(elo: int) -> int: + return max(_UCI_ELO_MIN, min(_UCI_ELO_MAX, elo)) + + +def sample_game_config( + openings: list, + game_id: int, + seed: int, + annotate_evals: bool = True, +) -> GameConfig: + """Randomly sample configuration for one game.""" + rng = random.Random(seed) + tier = rng.choices(_STRENGTH_TIERS, weights=_TIER_WEIGHTS, k=1)[0] + _, welo, belo, wdepth, bdepth = tier + eco, name, fen, moves = rng.choice(openings) + temperature = rng.uniform(60.0, 150.0) + temp_cutoff = rng.randint(8, 20) + return GameConfig( + white_elo=welo, + black_elo=belo, + white_depth=wdepth, + black_depth=bdepth, + opening_eco=eco, + opening_name=name, + opening_fen=fen, + opening_moves=list(moves), + temperature=temperature, + temp_cutoff_ply=temp_cutoff, + annotate_evals=annotate_evals, + game_id=game_id, + seed=seed, + ) + + +def temperature_sample( + multipv_info: list, + temperature_cp: float, + board: chess.Board, + rng: random.Random, +) -> chess.Move: + """ + Sample a move from MultiPV engine results using softmax temperature. + + Args: + multipv_info: List of info dicts from engine.analyse(..., multipv=N). + temperature_cp: Temperature in centipawns. Higher = more random. + board: Current board (used as fallback for legal moves). + rng: Seeded RNG for reproducibility. + + Returns: + A chess.Move sampled proportionally to exp(score / temperature_cp). + """ + moves: list[chess.Move] = [] + raw_scores: list[float] = [] + + for info in multipv_info: + if "pv" not in info or not info["pv"]: + continue + move = info["pv"][0] + score_obj = info["score"].relative + # Convert mate scores: mate in N → large finite value + cp = score_obj.score(mate_score=2000) + if cp is None: + continue + moves.append(move) + raw_scores.append(float(cp) / temperature_cp) + + if not moves: + legal = list(board.legal_moves) + return rng.choice(legal) + + # Numerically stable softmax + arr = np.array(raw_scores) + arr -= arr.max() + weights = np.exp(arr) + weights /= weights.sum() + + idx = np.random.choice(len(moves), p=weights) + return moves[idx] + + +def generate_game(sf_path: str, config: GameConfig) -> chess.pgn.Game: + """ + Generate one complete Stockfish self-play game. + + Returns a chess.pgn.Game object ready for PGN serialisation. + Raises on engine errors; caller should catch and log. + """ + rng = random.Random(config.seed) + np.random.seed(config.seed) + + engine = chess.engine.SimpleEngine.popen_uci(sf_path) + try: + board = chess.Board(config.opening_fen) + + # Apply opening moves + for uci in config.opening_moves: + mv = chess.Move.from_uci(uci) + if mv not in board.legal_moves: + break # opening line invalid for this position; stop early + board.push(mv) + + # Build PGN game object and replay the opening into it + game = chess.pgn.Game() + game.setup(chess.Board(config.opening_fen)) + node: chess.pgn.GameNode = game + for mv in board.move_stack: + node = node.add_variation(mv) + + # PGN headers + today = date.today().strftime("%Y.%m.%d") + game.headers["Event"] = "SF-SelfPlay" + game.headers["Site"] = "local" + game.headers["Date"] = today + game.headers["Round"] = str(config.game_id + 1) + game.headers["White"] = f"Stockfish_{config.white_elo}" + game.headers["Black"] = f"Stockfish_{config.black_elo}" + game.headers["Result"] = "*" + game.headers["WhiteElo"] = str(config.white_elo) + game.headers["BlackElo"] = str(config.black_elo) + game.headers["ECO"] = config.opening_eco + game.headers["Opening"] = config.opening_name + game.headers["TimeControl"] = "-" + game.headers["SFDepthWhite"] = str(config.white_depth) + game.headers["SFDepthBlack"] = str(config.black_depth) + + ply = board.ply() + + while not board.is_game_over(claim_draw=True): + is_white = board.turn == chess.WHITE + depth = config.white_depth if is_white else config.black_depth + elo = config.white_elo if is_white else config.black_elo + + clamped_elo = _clamp_elo(elo) + engine.configure({ + "UCI_LimitStrength": True, + "UCI_Elo": clamped_elo, + }) + + eval_score: int | None = None + + if ply < config.temp_cutoff_ply and config.temperature > 0: + # Temperature sampling phase for opening diversity + multipv = min(5, board.legal_moves.count()) + if multipv < 1: + break + info_list = engine.analyse( + board, + chess.engine.Limit(depth=depth), + multipv=multipv, + ) + move = temperature_sample(info_list, config.temperature, board, rng) + else: + # Pure engine play + if config.annotate_evals: + info = engine.analyse(board, chess.engine.Limit(depth=depth)) + move = info["pv"][0] if info.get("pv") else None + eval_score = info["score"].white().score(mate_score=10000) + else: + result = engine.play(board, chess.engine.Limit(depth=depth)) + move = result.move + + if move is None or move not in board.legal_moves: + # Fallback: let engine pick freely + result = engine.play(board, chess.engine.Limit(depth=depth)) + move = result.move + + node = node.add_variation(move) + if eval_score is not None: + node.comment = f"[%eval {eval_score / 100:.2f}]" + + board.push(move) + ply += 1 + + # Finalise result + outcome = board.outcome(claim_draw=True) + result_str = outcome.result() if outcome else "*" + game.headers["Result"] = result_str + game.headers["Termination"] = ( + outcome.termination.name.replace("_", " ").title() + if outcome else "Unknown" + ) + + return game + + finally: + engine.quit() + + +# --------------------------------------------------------------------------- +# Worker entry point (used by multiprocessing.Pool) +# --------------------------------------------------------------------------- + +def _worker(args: tuple) -> str | None: + """Generate one game and return its PGN string, or None on failure.""" + sf_path, config = args + try: + game = generate_game(sf_path, config) + return str(game) + except Exception as exc: # noqa: BLE001 + print(f"[worker] game {config.game_id} failed: {exc}", file=sys.stderr) + return None + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main() -> None: + parser = argparse.ArgumentParser( + description="Generate Stockfish self-play PGN training data.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--sf-path", + default="./stockfish/stockfish-ubuntu-x86-64-avx2", + help="Path to Stockfish binary.", + ) + parser.add_argument( + "--output", + default="selfplay_data/selfplay.pgn", + help="Output PGN file path.", + ) + parser.add_argument( + "--num-games", type=int, default=100, + help="Number of games to generate.", + ) + parser.add_argument( + "--seed", type=int, default=42, + help="Base random seed (each game gets seed = base + game_id).", + ) + parser.add_argument( + "--workers", type=int, default=1, + help="Number of parallel worker processes (1 = sequential).", + ) + parser.add_argument( + "--no-evals", action="store_true", + help="Skip %eval annotations (faster generation).", + ) + args = parser.parse_args() + + if not os.path.isfile(args.sf_path): + print(f"[error] Stockfish binary not found: {args.sf_path}", file=sys.stderr) + sys.exit(1) + + os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True) + + openings = load_openings() + configs = [ + sample_game_config( + openings, + game_id=i, + seed=args.seed + i, + annotate_evals=not args.no_evals, + ) + for i in range(args.num_games) + ] + + t_start = time.monotonic() + done = 0 + failed = 0 + + with open(args.output, "w") as out_file: + if args.workers <= 1: + # --- Sequential --- + for cfg in configs: + pgn_str = _worker((args.sf_path, cfg)) + if pgn_str: + out_file.write(pgn_str + "\n\n") + out_file.flush() + done += 1 + else: + failed += 1 + elapsed = time.monotonic() - t_start + rate = done / elapsed if elapsed > 0 else 0 + print( + f"\r[{done+failed}/{args.num_games}] " + f"done={done} failed={failed} " + f"rate={rate:.1f} games/s", + end="", + flush=True, + ) + else: + # --- Parallel --- + work_items = [(args.sf_path, cfg) for cfg in configs] + with Pool(processes=args.workers) as pool: + for pgn_str in pool.imap_unordered(_worker, work_items): + if pgn_str: + out_file.write(pgn_str + "\n\n") + out_file.flush() + done += 1 + else: + failed += 1 + elapsed = time.monotonic() - t_start + rate = done / elapsed if elapsed > 0 else 0 + print( + f"\r[{done+failed}/{args.num_games}] " + f"done={done} failed={failed} " + f"rate={rate:.1f} games/s", + end="", + flush=True, + ) + + elapsed = time.monotonic() - t_start + print( + f"\n\nFinished: {done} games saved to {args.output} " + f"({failed} failed) in {elapsed:.1f}s " + f"({done/elapsed:.1f} games/s)" + ) + + +if __name__ == "__main__": + main() diff --git a/grok_tracker.py b/grok_tracker.py new file mode 100644 index 0000000..91e3d2b --- /dev/null +++ b/grok_tracker.py @@ -0,0 +1,169 @@ +"""Grokking detection and Grokfast acceleration for training. + +GrokTracker monitors weight norms, gradient norms, effective rank of +activations, and train/test loss gap to detect the grokking transition +(delayed generalization after memorization). + +gradfilter_ema implements the Grokfast EMA gradient filter (Lee et al., 2024) +which amplifies slow-varying gradient components, achieving >50x speedup +of grokking with negligible computational overhead. +""" + +from __future__ import annotations + +import torch +from torch import Tensor, nn + + +class GrokTracker: + """Tracks grokking-related metrics during training. + + Monitors: + - Weight L2 norms (total) — rising = memorization, falling = generalization + - Gradient L2 norms — spikes often precede grokking transition + - Effective rank of activations (via SVD entropy) — drop = found structure + - Train/test loss gap — sustained gap + sudden test improvement = grokking + - Grokking onset detection — sustained weight norm decrease + """ + + def __init__(self, model: nn.Module, log_path: str | None = None) -> None: + self.model = model + self.weight_norm_history: list[float] = [] + self.grad_norm_history: list[float] = [] + self.train_loss_history: list[float] = [] + self.test_loss_history: list[float] = [] + self.log_path = log_path + self._log_file = open(log_path, "w") if log_path else None + + def compute_weight_norms(self) -> float: + """Total L2 norm of all parameters.""" + total_sq = sum( + p.data.norm(2).item() ** 2 for p in self.model.parameters() + ) + norm = total_sq**0.5 + self.weight_norm_history.append(norm) + return norm + + def compute_gradient_norms(self) -> float: + """Total L2 norm of all gradients (call after backward).""" + total_sq = sum( + p.grad.data.norm(2).item() ** 2 + for p in self.model.parameters() + if p.grad is not None + ) + norm = total_sq**0.5 + self.grad_norm_history.append(norm) + return norm + + def compute_effective_rank(self, activations: Tensor) -> float: + """Effective dimensionality via SVD entropy. + + Low rank = model found structured, low-dimensional representation. + High rank = representations are spread/random (memorization). + + Args: + activations: Tensor of shape [..., d_model] from a hidden layer. + + Returns: + Effective rank (1.0 = degenerate, d_model = fully random). + """ + flat = activations.reshape(-1, activations.shape[-1]).float() + svs = torch.linalg.svdvals(flat) + p = svs / svs.sum() + p = p[p > 1e-12] + entropy = -(p * p.log()).sum() + return torch.exp(entropy).item() + + def check_grokking_onset(self, window: int = 50) -> bool: + """Detect sustained weight norm decrease (cleanup phase signal). + + Returns True if average weight norm over the last `window` entries + is <95% of the preceding window — indicates the model is pruning + memorization circuits. + """ + if len(self.weight_norm_history) < 2 * window: + return False + recent = sum(self.weight_norm_history[-window:]) / window + earlier = sum(self.weight_norm_history[-2 * window : -window]) / window + return recent < 0.95 * earlier + + def log_epoch( + self, epoch: int, train_loss: float, test_loss: float + ) -> dict[str, float]: + """Log epoch-level metrics and return them as a dict.""" + self.train_loss_history.append(train_loss) + self.test_loss_history.append(test_loss) + w_norm = self.compute_weight_norms() + gap = test_loss - train_loss # positive = overfitting + grok = self.check_grokking_onset() + + line = ( + f"[grok] Epoch {epoch} | " + f"W_norm: {w_norm:.2f} | " + f"Train: {train_loss:.4f} | Test: {test_loss:.4f} | " + f"Gap: {gap:+.4f} | " + f"Grokking: {'>>> YES <<<' if grok else 'no'}" + ) + print(line) + if self._log_file: + self._log_file.write(line + "\n") + self._log_file.flush() + + return { + "weight_norm": w_norm, + "loss_gap": gap, + "grokking_onset": grok, + } + + def log_batch(self, step: int) -> dict[str, float]: + """Log batch-level metrics (lightweight — weight + grad norms only). + + Call after backward() but before optimizer.step(). + """ + w_norm = self.compute_weight_norms() + g_norm = self.compute_gradient_norms() + return {"weight_norm": w_norm, "gradient_norm": g_norm} + + def close(self) -> None: + if self._log_file: + self._log_file.close() + self._log_file = None + + +def gradfilter_ema( + model: nn.Module, + grads: dict[str, Tensor] | None = None, + alpha: float = 0.98, + lamb: float = 2.0, +) -> dict[str, Tensor]: + """Grokfast EMA gradient filter (Lee et al., 2024). + + Amplifies slow-varying gradient components by maintaining an EMA + of gradients and adding the amplified EMA back to current gradients. + Achieves >50x speedup of grokking with negligible overhead. + + Call after loss.backward() and before optimizer.step(): + loss.backward() + ema_grads = gradfilter_ema(model, ema_grads) + optimizer.step() + + Args: + model: Model whose gradients to filter. + grads: Previous EMA state (None on first call — initializes). + alpha: EMA decay factor (0.98 = slow, captures long-term patterns). + lamb: Amplification factor (2.0 = double the slow components). + + Returns: + Updated EMA state dict (pass to next call). + """ + if grads is None: + return { + n: p.grad.data.detach().clone() + for n, p in model.named_parameters() + if p.requires_grad and p.grad is not None + } + for n, p in model.named_parameters(): + if p.requires_grad and p.grad is not None: + grads[n] = grads[n] * alpha + p.grad.data.detach() * (1 - alpha) + p.grad.data += grads[n] * lamb + return grads diff --git a/inference_test.py b/inference_test.py index 9720a6c..07179e0 100644 --- a/inference_test.py +++ b/inference_test.py @@ -1,35 +1,34 @@ import torch import chess -from chess_loader import ChessDataset -from chessformer import ChessTransformer +import sys +import chessformer +sys.modules['transformer'] = chessformer +from chessformer import ChessTransformer, ChessTransformerV2 from chess_moves_to_input_data import get_board_str, switch_player, switch_move +from model_utils import ( + detect_device, + load_model, + preprocess_board as preprocess_v2, + preprocess_board_v1 as preprocess, +) +from policy import greedy_move_v2 from torch.utils.data import DataLoader from copy import deepcopy import time # Configuration -MODEL = "2000_elo_pos_engine_best_test_whole.pth" +MODEL = "2000_elo_pos_engine.pth" -# Model and device setup -if torch.backends.mps.is_available(): - device = torch.device("mps") -elif torch.cuda.is_available(): - device = torch.device("cuda") -else: - device = torch.device("cpu") +# Changed: --device flag to override auto-detection (auto/cuda/mps/cpu) +import argparse +_parser = argparse.ArgumentParser() +_parser.add_argument('--device', type=str, default='auto', choices=['auto', 'cuda', 'mps', 'cpu']) +_args = _parser.parse_args() -model = torch.load(f'models/{MODEL}').to(device) +device = detect_device(_args.device) -# Preprocessing function -def preprocess(board): - """ - Converts a chess board state to a tensor representation. - """ - board_str = get_board_str(board, white_side=board.turn) - piece_to_index = {'.': 1, 'P': 2, 'N': 3, 'B': 4, 'R': 5, 'Q': 6, 'K': 7, - 'p': 8, 'n': 9, 'b': 10, 'r': 11, 'q': 12, 'k': 13} - board_pieces = [piece_to_index[p] for p in board_str] - return torch.tensor([board_pieces], dtype=torch.long).to(device) +# Auto-detect V1/V2 from checkpoint format +model, _model_version, _model_cfg = load_model(f'models/{MODEL}', device) # Helper functions for postprocessing def sq_to_str(sq): @@ -76,10 +75,17 @@ def postprocess_valid(output, board: chess.Board, rep_mv=""): count = 0 while not (board.is_checkmate() or board.is_stalemate() or board.is_insufficient_material() or board.can_claim_draw()): - input_tensors = preprocess(board) count += 1 with torch.no_grad(): - output = model(*input_tensors) - uci_move = postprocess_valid(output, board) + if _model_version == 'v2': + board_t, feat_t = preprocess_v2(board, device) + policy_logits, promo_logits, wdl, ply = model(board_t, feat_t) + move = greedy_move_v2(board, policy_logits[0], promo_logits[0]) + uci_move = move.uci() + print(f'WDL: W={wdl[0,0]:.2f} D={wdl[0,1]:.2f} L={wdl[0,2]:.2f}') + else: + input_tensors = preprocess(board, device) + output = model(input_tensors) + uci_move = postprocess_valid(output, board) board.push(chess.Move.from_uci(uci_move)) print(f'Predicted {count}\n', board) diff --git a/mcts.py b/mcts.py new file mode 100644 index 0000000..1b6d600 --- /dev/null +++ b/mcts.py @@ -0,0 +1,250 @@ +"""Monte Carlo Tree Search for ChessTransformerV2. + +Inspired by alpha-zero-general (suragnair/alpha-zero-general), but works +directly with chess.Board and our model — no generic Game/NNet adapters needed. + +Uses virtual-loss batching: collects multiple leaves per batch, evaluates them +in a single forward pass (~5-7x faster than one-at-a-time with 25+ sims). + +Usage: + mcts = MCTS(model, device, num_simulations=50, cpuct=1.25) + visit_policy, root_wdl = mcts.search(board) + # visit_policy: dict[chess.Move, float] — normalized visit counts + # root_wdl: (win, draw, loss) probabilities from root forward pass +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass, field + +import chess +import numpy as np +import torch + +from model_utils import preprocess_board, preprocess_boards_batch +from policy import legal_move_policy_v2 + + +@dataclass +class MCTSNode: + """Single node in the MCTS search tree.""" + + board: chess.Board + parent: MCTSNode | None = None + move: chess.Move | None = None + children: dict[chess.Move, MCTSNode] = field(default_factory=dict) + visit_count: int = 0 + value_sum: float = 0.0 + prior: float = 0.0 + + def q_value(self) -> float: + return self.value_sum / self.visit_count if self.visit_count > 0 else 0.0 + + def is_expanded(self) -> bool: + return len(self.children) > 0 + + +class MCTS: + """AlphaZero-style MCTS using ChessTransformerV2 for priors and value. + + Virtual-loss batching: instead of evaluating one leaf at a time, + collects `batch_size` leaves per round using virtual losses to diversify + selection, then evaluates all in one forward pass. Cuts GPU round-trips + from num_simulations down to num_simulations / batch_size. + """ + + def __init__( + self, + model: torch.nn.Module, + device: torch.device, + num_simulations: int = 50, + cpuct: float = 1.25, + use_diffusion: bool = False, + batch_size: int = 8, + dirichlet_alpha: float = 0.3, + dirichlet_epsilon: float = 0.25, + ): + self.model = model + self.device = device + self.num_simulations = num_simulations + self.cpuct = cpuct + self.use_diffusion = use_diffusion + self.batch_size = batch_size + self.dirichlet_alpha = dirichlet_alpha + self.dirichlet_epsilon = dirichlet_epsilon + self._last_wdl: tuple[float, float, float] = (0.0, 0.5, 0.5) + + def search( + self, board: chess.Board + ) -> tuple[dict[chess.Move, float], tuple[float, float, float]]: + """Run MCTS from position. + + Returns: + (visit_policy, root_wdl) where visit_policy maps legal moves + to normalized visit fractions, and root_wdl is (W, D, L) from + the root node's neural network evaluation. + """ + root = MCTSNode(board=board.copy()) + self._expand_and_evaluate(root) + root_wdl_tuple = self._last_wdl + + # Dirichlet noise at root — forces exploration of unexpected moves + if root.children and self.dirichlet_epsilon > 0: + noise = np.random.dirichlet( + [self.dirichlet_alpha] * len(root.children) + ) + for child, eta in zip(root.children.values(), noise): + child.prior = ( + (1 - self.dirichlet_epsilon) * child.prior + + self.dirichlet_epsilon * eta + ) + + remaining = self.num_simulations + while remaining > 0: + batch = min(self.batch_size, remaining) + + # SELECT: collect leaves, applying virtual loss to diversify + leaves: list[MCTSNode] = [] + for _ in range(batch): + node = root + while node.is_expanded() and not node.board.is_game_over(): + node = self._select_child(node) + # Virtual loss: discourage other sims from picking same node + node.visit_count += 1 + node.value_sum -= 1.0 + leaves.append(node) + + # EXPAND + EVALUATE (batched) + value_by_id: dict[int, float] = {} + to_expand: list[MCTSNode] = [] + seen: set[int] = set() + + for node in leaves: + nid = id(node) + if nid in seen: + continue + seen.add(nid) + if node.board.is_game_over(): + outcome = node.board.outcome() + value_by_id[nid] = ( + 0.0 if (outcome is None or outcome.winner is None) else -1.0 + ) + elif node.is_expanded(): + # Already expanded (by previous batch) + value_by_id[nid] = node.q_value() + else: + to_expand.append(node) + + if to_expand: + values = self._batch_expand_and_evaluate(to_expand) + for node, val in zip(to_expand, values): + value_by_id[id(node)] = val + + # UNDO virtual loss + BACKUP + for node in leaves: + node.visit_count -= 1 + node.value_sum += 1.0 + self._backup(node, value_by_id[id(node)]) + + remaining -= batch + + total = sum(c.visit_count for c in root.children.values()) + if total == 0: + return {}, root_wdl_tuple + policy = { + move: child.visit_count / total + for move, child in root.children.items() + } + return policy, root_wdl_tuple + + def _select_child(self, node: MCTSNode) -> MCTSNode: + """Pick child with highest UCB score.""" + total_visits = sum(c.visit_count for c in node.children.values()) + sqrt_total = math.sqrt(total_visits) if total_visits > 0 else 1.0 + + best_score = -math.inf + best_child = None + for child in node.children.values(): + q = child.q_value() + exploration = ( + self.cpuct * child.prior * sqrt_total / (1 + child.visit_count) + ) + score = q + exploration + if score > best_score: + best_score = score + best_child = child + return best_child # type: ignore[return-value] + + def _expand_node( + self, + node: MCTSNode, + policy_logits: torch.Tensor, + promo_logits: torch.Tensor, + ) -> None: + """Create child nodes with neural network priors.""" + moves, probs, _log_probs = legal_move_policy_v2( + node.board, policy_logits, promo_logits + ) + for move, prior in zip(moves, probs): + child_board = node.board.copy() + child_board.push(move) + node.children[move] = MCTSNode( + board=child_board, + parent=node, + move=move, + prior=prior.item(), + ) + + @torch.no_grad() + def _expand_and_evaluate(self, node: MCTSNode) -> float: + """Expand single leaf node (used for root). Returns value W - L.""" + if node.board.is_game_over(): + outcome = node.board.outcome() + if outcome is None or outcome.winner is None: + self._last_wdl = (0.0, 1.0, 0.0) + return 0.0 + self._last_wdl = (0.0, 0.0, 1.0) + return -1.0 + + board_t, feat_t = preprocess_board(node.board, self.device) + policy_logits, promo_logits, wdl, _ply = self.model( + board_t, feat_t, use_diffusion=self.use_diffusion + ) + + w, d, l = wdl[0, 0].item(), wdl[0, 1].item(), wdl[0, 2].item() + self._last_wdl = (w, d, l) + self._expand_node(node, policy_logits[0], promo_logits[0]) + return w - l + + @torch.no_grad() + def _batch_expand_and_evaluate( + self, nodes: list[MCTSNode] + ) -> list[float]: + """Expand multiple leaf nodes in a single batched forward pass.""" + if len(nodes) == 1: + return [self._expand_and_evaluate(nodes[0])] + + boards = [node.board for node in nodes] + boards_t, feats_t = preprocess_boards_batch(boards, self.device) + + policy_logits, promo_logits, wdl, _ply = self.model( + boards_t, feats_t, use_diffusion=self.use_diffusion + ) + + values = [] + for i, node in enumerate(nodes): + w, l = wdl[i, 0].item(), wdl[i, 2].item() + self._expand_node(node, policy_logits[i], promo_logits[i]) + values.append(w - l) + + return values + + @staticmethod + def _backup(node: MCTSNode, value: float) -> None: + """Propagate value up the tree, negating at each level.""" + while node is not None: + node.visit_count += 1 + node.value_sum += value + value = -value + node = node.parent # type: ignore[assignment] diff --git a/model_utils.py b/model_utils.py new file mode 100644 index 0000000..eb6fd89 --- /dev/null +++ b/model_utils.py @@ -0,0 +1,178 @@ +"""Shared utilities for model loading, device detection, preprocessing, and loss. + +Consolidates duplicated logic from train_model.py, inference_test.py, +play_gui.py, and selfplay_loop.py into a single module. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import chess +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +import chessformer +from chess_loader import PIECE_TO_INDEX, compute_features +from chess_moves_to_input_data import get_board_str +from chessformer import ChessTransformer, ChessTransformerV2 + +# Allow unpickling old models saved under the name 'transformer' +sys.modules["transformer"] = chessformer + + +def detect_device(preference: str = "auto") -> torch.device: + """Resolve device string to a torch.device. + + Args: + preference: "auto", "cuda", "mps", or "cpu". + + Returns: + Resolved torch.device. + """ + if preference != "auto": + return torch.device(preference) + if torch.cuda.is_available(): + return torch.device("cuda") + if torch.backends.mps.is_available(): + return torch.device("mps") + return torch.device("cpu") + + +def find_latest_model(directory: str = "models", version: str = "v2") -> str | None: + """Find the most recently modified V2 model in directory. + + Returns path string or None if no models found. + """ + models_dir = Path(directory) + if not models_dir.is_dir(): + return None + candidates = sorted(models_dir.glob("*_v2.pth"), key=lambda p: p.stat().st_mtime) + if not candidates: + # fallback: any .pth file + candidates = sorted(models_dir.glob("*.pth"), key=lambda p: p.stat().st_mtime) + return str(candidates[-1]) if candidates else None + + +def load_model( + path: str, device: torch.device +) -> tuple[nn.Module, str, dict | None]: + """Load a V1 or V2 model checkpoint. + + Handles: + - V2 checkpoints: {'version': 'v2', 'state_dict': ..., 'config': ...} + - V1 state dicts (including RL adapter prefix stripping) + - V1 pickled model objects + + Args: + path: Path to .pth checkpoint. + device: Target device. + + Returns: + (model, version, config) where version is 'v1' or 'v2', + and config is the V2 config dict (None for V1). + """ + obj = torch.load(path, weights_only=False, map_location="cpu") + + if isinstance(obj, dict) and obj.get("version") == "v2": + cfg = obj["config"] + model = ChessTransformerV2(**cfg, dropout=0.0) + model.load_state_dict(obj["state_dict"]) + return model.to(device), "v2", cfg + + if isinstance(obj, dict): + sd = obj + # RL checkpoint: keys prefixed with "adapter.transformer." + prefix = "adapter.transformer." + if any(k.startswith(prefix) for k in sd): + sd = {k[len(prefix):]: v for k, v in sd.items() if k.startswith(prefix)} + # Infer V1 architecture from state dict + if "embedding.weight" not in sd: + raise ValueError(f"Unrecognised checkpoint format in {path}") + inp_dict = sd["embedding.weight"].shape[0] + out_dict = sd["linear_output.weight"].shape[0] + d_model = sd["embedding.weight"].shape[1] + d_hid = sd["transformer_encoder.layers.0.linear1.weight"].shape[0] + nlayers = sum(1 for k in sd if k.endswith(".self_attn.in_proj_weight")) + nhead = max(1, d_model // 64) + model = ChessTransformer( + inp_dict, out_dict, d_model, nhead, d_hid, nlayers, dropout=0.0 + ) + model.load_state_dict(sd) + return model.to(device), "v1", None + + # Pickled model object (legacy V1) + if isinstance(obj, nn.Module): + return obj.to(device), "v1", None + + raise ValueError(f"Cannot load model from {path}") + + +def preprocess_board( + board: chess.Board, device: torch.device +) -> tuple[Tensor, Tensor]: + """Convert chess.Board to V2 model inputs. + + Board is always oriented for the current player (flipped if black to move). + + Returns: + (board_tensor[1,64], features_tensor[1,14]) + """ + board_str = get_board_str(board, white_side=board.turn) + board_pieces = [PIECE_TO_INDEX[p] for p in board_str] + features = compute_features(board_str) + return ( + torch.tensor([board_pieces], dtype=torch.long, device=device), + torch.tensor([features], dtype=torch.float32, device=device), + ) + + +def preprocess_boards_batch( + boards: list[chess.Board], device: torch.device +) -> tuple[Tensor, Tensor]: + """Batch version of preprocess_board for multiple boards at once. + + Builds a single [N, 64] int tensor and [N, 14] float tensor from N boards, + using one torch.tensor() call each instead of N separate allocations. + """ + pieces_list = [] + feats_list = [] + for board in boards: + board_str = get_board_str(board, white_side=board.turn) + pieces_list.append([PIECE_TO_INDEX[p] for p in board_str]) + feats_list.append(compute_features(board_str)) + return ( + torch.tensor(pieces_list, dtype=torch.long, device=device), + torch.tensor(feats_list, dtype=torch.float32, device=device), + ) + + +def preprocess_board_v1( + board: chess.Board, device: torch.device +) -> Tensor: + """Convert chess.Board to V1 model input (board only, no features). + + Returns: + board_tensor[1,64] + """ + board_str = get_board_str(board, white_side=board.turn) + board_pieces = [PIECE_TO_INDEX[p] for p in board_str] + return torch.tensor([board_pieces], dtype=torch.long, device=device) + + +def compute_loss_v2( + model: ChessTransformerV2, + boards: Tensor, + features: Tensor, + from_sq: Tensor, + to_sq: Tensor, + wdl_target: Tensor, +) -> Tensor: + """V2 multi-task loss: policy cross-entropy + 0.5 * WDL cross-entropy.""" + policy_logits, _promo, wdl_pred, _ply = model(boards, features) + B = boards.shape[0] + policy_loss = F.cross_entropy(policy_logits.reshape(B, -1), from_sq * 64 + to_sq) + wdl_loss = -(wdl_target * torch.log(wdl_pred + 1e-8)).sum(dim=-1).mean() + return policy_loss + 0.5 * wdl_loss diff --git a/models/2000_elo_pos_engine.pth b/models/2000_elo_pos_engine.pth new file mode 100644 index 0000000..b6d531f --- /dev/null +++ b/models/2000_elo_pos_engine.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:545f53bc7a9fcf8b9b6af104817ec620055129d32f80c79dc18fb32e9c91adb2 +size 111317943 diff --git a/models/2500_elo_pos_engine_v2.pth b/models/2500_elo_pos_engine_v2.pth new file mode 100644 index 0000000..dafa5cb --- /dev/null +++ b/models/2500_elo_pos_engine_v2.pth @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0f432c27dec3671d93abc4cbd95ddc85d2290824b2212cd00c6a50ae2a75ba0c +size 169998229 diff --git a/noise_schedule.py b/noise_schedule.py new file mode 100644 index 0000000..bbf4f2b --- /dev/null +++ b/noise_schedule.py @@ -0,0 +1,94 @@ +"""noise_schedule.py — Cosine noise schedule for continuous DDPM. + +Implements the forward diffusion process: gradually adding Gaussian noise +to clean latent states over T timesteps. The reverse process (denoising) +is handled by the DiT model in diffusion_model.py. + +Key formula: + x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * noise + +where alpha_bar_t follows a cosine curve from ~1 (clean) to ~0 (pure noise). + +Reference: Nichol & Dhariwal "Improved DDPM" (2021), Section 3.2. +""" + +import math +import torch +from torch import Tensor + + +class CosineNoiseSchedule: + """Cosine noise schedule with precomputed diffusion coefficients. + + Args: + T: number of diffusion timesteps (20 for chess, per DiffuSearch finding) + s: small offset to prevent alpha_bar from being exactly 1 at t=0 + """ + + def __init__(self, T: int = 20, s: float = 0.008) -> None: + self.T = T + + # alpha_bar: cumulative signal retention, shape [T+1] + # alpha_bar[0] ≈ 1 (clean), alpha_bar[T] ≈ 0 (pure noise) + steps = torch.arange(T + 1, dtype=torch.float64) + f_t = torch.cos(((steps / T) + s) / (1 + s) * (math.pi / 2)) ** 2 + alpha_bar = f_t / f_t[0] + # Clamp to avoid numerical issues at boundaries + alpha_bar = alpha_bar.clamp(min=1e-5, max=1.0 - 1e-5) + + self.alpha_bar = alpha_bar.float() + # sqrt values used in q_sample + self.sqrt_alpha_bar = self.alpha_bar.sqrt() + self.sqrt_one_minus_alpha_bar = (1.0 - self.alpha_bar).sqrt() + + # Per-step alpha and beta (needed for reverse process / denoising step) + # alpha_t = alpha_bar_t / alpha_bar_{t-1} + # beta_t = 1 - alpha_t + alpha = torch.cat([ + torch.tensor([1.0]), + self.alpha_bar[1:] / self.alpha_bar[:-1], + ]) + alpha = alpha.clamp(min=1e-5, max=1.0) + self.beta = (1.0 - alpha).float() + self.alpha = alpha.float() + + # Posterior variance for denoising: beta_tilde_t = beta_t * (1 - alpha_bar_{t-1}) / (1 - alpha_bar_t) + alpha_bar_prev = torch.cat([torch.tensor([1.0]), self.alpha_bar[:-1]]) + self.posterior_variance = ( + self.beta * (1.0 - alpha_bar_prev) / (1.0 - self.alpha_bar) + ).clamp(min=1e-20) + + def q_sample(self, x_0: Tensor, t: Tensor, noise: Tensor | None = None) -> Tensor: + """Forward process: add noise to clean data. + + Args: + x_0: clean data, any shape [B, ...] + t: timestep indices, shape [B], values in [0, T] + noise: optional pre-sampled Gaussian noise, same shape as x_0 + + Returns: + x_t: noisy data at timestep t, same shape as x_0 + """ + if noise is None: + noise = torch.randn_like(x_0) + + # Gather coefficients for each batch element, reshape for broadcasting + sqrt_ab = self.sqrt_alpha_bar[t] + sqrt_1_ab = self.sqrt_one_minus_alpha_bar[t] + + # Reshape to broadcast: [B] -> [B, 1, 1, ...] matching x_0 dims + while sqrt_ab.dim() < x_0.dim(): + sqrt_ab = sqrt_ab.unsqueeze(-1) + sqrt_1_ab = sqrt_1_ab.unsqueeze(-1) + + return sqrt_ab * x_0 + sqrt_1_ab * noise + + def to(self, device: torch.device) -> "CosineNoiseSchedule": + """Move all precomputed tensors to a device.""" + self.alpha_bar = self.alpha_bar.to(device) + self.sqrt_alpha_bar = self.sqrt_alpha_bar.to(device) + self.sqrt_one_minus_alpha_bar = self.sqrt_one_minus_alpha_bar.to(device) + self.beta = self.beta.to(device) + self.alpha = self.alpha.to(device) + self.posterior_variance = self.posterior_variance.to(device) + return self diff --git a/openings.py b/openings.py new file mode 100644 index 0000000..1bfb0f7 --- /dev/null +++ b/openings.py @@ -0,0 +1,289 @@ +""" +Opening book for self-play data generation. + +Each entry is a tuple: (eco_code, name, fen, moves_uci) +- eco_code: ECO classification string +- name: Human-readable opening name +- fen: Starting FEN (always the standard starting position here; + moves_uci is applied on top of it) +- moves_uci: List of UCI moves to reach the opening position + +The standard starting FEN is used as the base for all openings. +moves_uci lists the moves that define the opening line; after these +are played the engine takes over with temperature sampling. +""" + +import random + +START_FEN = "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1" + +# fmt: off +OPENINGS: list[tuple[str, str, str, list[str]]] = [ + # ------------------------------------------------------------------ # + # Open Games (1. e4 e5) + # ------------------------------------------------------------------ # + ("C20", "King's Pawn Game", START_FEN, ["e2e4", "e7e5"]), + ("C21", "Center Game", START_FEN, ["e2e4", "e7e5", "d2d4", "e5d4"]), + ("C23", "Bishop's Opening", START_FEN, ["e2e4", "e7e5", "f1c4"]), + ("C25", "Vienna Game", START_FEN, ["e2e4", "e7e5", "b1c3"]), + ("C30", "King's Gambit", START_FEN, ["e2e4", "e7e5", "f2f4"]), + ("C31", "King's Gambit Declined", START_FEN, ["e2e4", "e7e5", "f2f4", "d7d5"]), + ("C33", "King's Gambit Accepted", START_FEN, ["e2e4", "e7e5", "f2f4", "e5f4"]), + ("C40", "Latvian Gambit", START_FEN, ["e2e4", "e7e5", "g1f3", "f7f5"]), + ("C41", "Philidor Defense", START_FEN, ["e2e4", "e7e5", "g1f3", "d7d6"]), + ("C42", "Petrov Defense", START_FEN, ["e2e4", "e7e5", "g1f3", "g8f6"]), + ("C44", "Scotch Game", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "d2d4"]), + ("C45", "Scotch Game: main line", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "d2d4", "e5d4", "f3d4"]), + ("C46", "Three Knights Game", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "b1c3"]), + ("C47", "Four Knights Game", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "b1c3", "g8f6"]), + ("C50", "Italian Game", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1c4"]), + ("C51", "Evans Gambit", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1c4", "f8c5", "b2b4"]), + ("C53", "Italian: Classical", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1c4", "f8c5", "c2c3"]), + ("C54", "Italian: Giuoco Piano", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1c4", "f8c5", "c2c3", "g8f6", "d2d4"]), + ("C55", "Two Knights Defense", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1c4", "g8f6"]), + ("C56", "Two Knights: Fried Liver", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1c4", "g8f6", "f3g5", "d7d5", "e4d5", "c6a5"]), + ("C60", "Ruy Lopez", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1b5"]), + ("C61", "Ruy Lopez: Bird Defense", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1b5", "c6d4"]), + ("C62", "Ruy Lopez: Steinitz Defense", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1b5", "d7d6"]), + ("C63", "Ruy Lopez: Schliemann", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1b5", "f7f5"]), + ("C65", "Ruy Lopez: Berlin Defense", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1b5", "g8f6"]), + ("C67", "Ruy Lopez: Berlin Endgame", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1b5", "g8f6", "e1g1", "f6e4", "d1e1", "e4d6", "f3e5", "c6e5"]), + ("C68", "Ruy Lopez: Exchange", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1b5", "a7a6", "b5c6"]), + ("C70", "Ruy Lopez: Morphy Defense", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1b5", "a7a6", "b5a4"]), + ("C78", "Ruy Lopez: Archangel", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1b5", "a7a6", "b5a4", "g8f6", "e1g1", "f8b4"]), + ("C80", "Ruy Lopez: Open", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1b5", "a7a6", "b5a4", "g8f6", "e1g1", "f6e4"]), + ("C84", "Ruy Lopez: Closed", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1b5", "a7a6", "b5a4", "g8f6", "e1g1", "f8e7"]), + ("C88", "Ruy Lopez: Anti-Marshall", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1b5", "a7a6", "b5a4", "g8f6", "e1g1", "f8e7", "f1e1", "b7b5", "a4b3", "e8g8"]), + ("C92", "Ruy Lopez: Marshall Attack", START_FEN, ["e2e4", "e7e5", "g1f3", "b8c6", "f1b5", "a7a6", "b5a4", "g8f6", "e1g1", "f8e7", "f1e1", "b7b5", "a4b3", "e8g8", "c2c3", "d7d5"]), + + # ------------------------------------------------------------------ # + # Sicilian Defense (1. e4 c5) + # ------------------------------------------------------------------ # + ("B20", "Sicilian Defense", START_FEN, ["e2e4", "c7c5"]), + ("B21", "Sicilian: Grand Prix Attack", START_FEN, ["e2e4", "c7c5", "b1c3", "b8c6", "f2f4"]), + ("B22", "Sicilian: Alapin", START_FEN, ["e2e4", "c7c5", "c2c3"]), + ("B23", "Sicilian: Closed", START_FEN, ["e2e4", "c7c5", "b1c3"]), + ("B27", "Sicilian: Hungarian", START_FEN, ["e2e4", "c7c5", "g1f3", "g7g6"]), + ("B30", "Sicilian: Old Sicilian", START_FEN, ["e2e4", "c7c5", "g1f3", "b8c6"]), + ("B32", "Sicilian: Löwenthal", START_FEN, ["e2e4", "c7c5", "g1f3", "b8c6", "d2d4", "c5d4", "f3d4", "e7e5"]), + ("B40", "Sicilian: Kan", START_FEN, ["e2e4", "c7c5", "g1f3", "e7e6"]), + ("B41", "Sicilian: Kan main", START_FEN, ["e2e4", "c7c5", "g1f3", "e7e6", "d2d4", "c5d4", "f3d4", "a7a6"]), + ("B44", "Sicilian: Taimanov", START_FEN, ["e2e4", "c7c5", "g1f3", "e7e6", "d2d4", "c5d4", "f3d4", "b8c6"]), + ("B50", "Sicilian: 2.Nf3 d6", START_FEN, ["e2e4", "c7c5", "g1f3", "d7d6"]), + ("B54", "Sicilian: Dragon", START_FEN, ["e2e4", "c7c5", "g1f3", "d7d6", "d2d4", "c5d4", "f3d4", "g8f6", "b1c3", "g7g6"]), + ("B57", "Sicilian: Classical", START_FEN, ["e2e4", "c7c5", "g1f3", "d7d6", "d2d4", "c5d4", "f3d4", "g8f6", "b1c3", "b8c6"]), + ("B58", "Sicilian: Boleslavsky", START_FEN, ["e2e4", "c7c5", "g1f3", "d7d6", "d2d4", "c5d4", "f3d4", "g8f6", "b1c3", "e7e5"]), + ("B60", "Sicilian: Richter-Rauzer", START_FEN, ["e2e4", "c7c5", "g1f3", "d7d6", "d2d4", "c5d4", "f3d4", "g8f6", "b1c3", "b8c6", "c1g5"]), + ("B70", "Sicilian: Dragon main", START_FEN, ["e2e4", "c7c5", "g1f3", "d7d6", "d2d4", "c5d4", "f3d4", "g8f6", "b1c3", "g7g6", "f1e2", "f8g7"]), + ("B72", "Sicilian: Scheveningen", START_FEN, ["e2e4", "c7c5", "g1f3", "d7d6", "d2d4", "c5d4", "f3d4", "g8f6", "b1c3", "e7e6"]), + ("B80", "Sicilian: Najdorf", START_FEN, ["e2e4", "c7c5", "g1f3", "d7d6", "d2d4", "c5d4", "f3d4", "g8f6", "b1c3", "a7a6"]), + ("B85", "Sicilian: Najdorf 6.Be2", START_FEN, ["e2e4", "c7c5", "g1f3", "d7d6", "d2d4", "c5d4", "f3d4", "g8f6", "b1c3", "a7a6", "f1e2", "e7e6"]), + ("B90", "Sicilian: Najdorf 6.Be3", START_FEN, ["e2e4", "c7c5", "g1f3", "d7d6", "d2d4", "c5d4", "f3d4", "g8f6", "b1c3", "a7a6", "c1e3"]), + ("B96", "Sicilian: Najdorf Poisoned Pawn", START_FEN, ["e2e4", "c7c5", "g1f3", "d7d6", "d2d4", "c5d4", "f3d4", "g8f6", "b1c3", "a7a6", "c1g5", "e7e6", "f2f4", "d8b6"]), + + # ------------------------------------------------------------------ # + # French Defense (1. e4 e6) + # ------------------------------------------------------------------ # + ("C00", "French Defense", START_FEN, ["e2e4", "e7e6"]), + ("C01", "French: Exchange", START_FEN, ["e2e4", "e7e6", "d2d4", "d7d5", "e4d5"]), + ("C02", "French: Advance", START_FEN, ["e2e4", "e7e6", "d2d4", "d7d5", "e4e5"]), + ("C03", "French: Tarrasch", START_FEN, ["e2e4", "e7e6", "d2d4", "d7d5", "b1d2"]), + ("C10", "French: Rubinstein", START_FEN, ["e2e4", "e7e6", "d2d4", "d7d5", "b1c3", "d5e4"]), + ("C11", "French: Classical", START_FEN, ["e2e4", "e7e6", "d2d4", "d7d5", "b1c3", "g8f6"]), + ("C14", "French: Classical Winawer", START_FEN, ["e2e4", "e7e6", "d2d4", "d7d5", "b1c3", "f8b4"]), + ("C18", "French: Winawer Poisoned Pawn", START_FEN, ["e2e4", "e7e6", "d2d4", "d7d5", "b1c3", "f8b4", "e4e5", "c7c5", "a2a3", "b4c3", "b2c3", "d8c7"]), + + # ------------------------------------------------------------------ # + # Caro-Kann (1. e4 c6) + # ------------------------------------------------------------------ # + ("B10", "Caro-Kann Defense", START_FEN, ["e2e4", "c7c6"]), + ("B12", "Caro-Kann: Advance", START_FEN, ["e2e4", "c7c6", "d2d4", "d7d5", "e4e5"]), + ("B13", "Caro-Kann: Exchange", START_FEN, ["e2e4", "c7c6", "d2d4", "d7d5", "e4d5"]), + ("B14", "Caro-Kann: Panov Attack", START_FEN, ["e2e4", "c7c6", "d2d4", "d7d5", "e4d5", "c6d5", "c2c4"]), + ("B15", "Caro-Kann: Classical", START_FEN, ["e2e4", "c7c6", "d2d4", "d7d5", "b1c3", "d5e4", "c3e4", "g8f6"]), + ("B17", "Caro-Kann: Steinitz", START_FEN, ["e2e4", "c7c6", "d2d4", "d7d5", "b1c3", "d5e4", "c3e4", "b8d7"]), + + # ------------------------------------------------------------------ # + # Pirc / Modern Defense + # ------------------------------------------------------------------ # + ("B06", "Modern Defense", START_FEN, ["e2e4", "g7g6"]), + ("B07", "Pirc Defense", START_FEN, ["e2e4", "d7d6", "d2d4", "g8f6", "b1c3", "g7g6"]), + ("B08", "Pirc: Classical", START_FEN, ["e2e4", "d7d6", "d2d4", "g8f6", "b1c3", "g7g6", "g1f3"]), + ("B09", "Pirc: Austrian Attack", START_FEN, ["e2e4", "d7d6", "d2d4", "g8f6", "b1c3", "g7g6", "f2f4"]), + + # ------------------------------------------------------------------ # + # Scandinavian (1. e4 d5) + # ------------------------------------------------------------------ # + ("B01", "Scandinavian Defense", START_FEN, ["e2e4", "d7d5", "e4d5", "d8d5", "b1c3", "d5a5"]), + ("B01b", "Scandinavian: 3...Qd6", START_FEN, ["e2e4", "d7d5", "e4d5", "d8d5", "b1c3", "d5d6"]), + + # ------------------------------------------------------------------ # + # Alekhine's Defense + # ------------------------------------------------------------------ # + ("B02", "Alekhine's Defense", START_FEN, ["e2e4", "g8f6"]), + ("B04", "Alekhine: Modern Variation", START_FEN, ["e2e4", "g8f6", "e4e5", "f6d5", "d2d4", "d7d6", "g1f3"]), + + # ------------------------------------------------------------------ # + # Queen's Gambit (1. d4 d5 2. c4) + # ------------------------------------------------------------------ # + ("D00", "Queen's Pawn Game", START_FEN, ["d2d4", "d7d5"]), + ("D01", "Richter-Veresov Attack", START_FEN, ["d2d4", "d7d5", "b1c3", "g8f6", "c1g5"]), + ("D02", "London System", START_FEN, ["d2d4", "d7d5", "g1f3", "g8f6", "c1f4"]), + ("D06", "Queen's Gambit", START_FEN, ["d2d4", "d7d5", "c2c4"]), + ("D07", "QGD: Chigorin Defense", START_FEN, ["d2d4", "d7d5", "c2c4", "b8c6"]), + ("D10", "QGD: Slav", START_FEN, ["d2d4", "d7d5", "c2c4", "c7c6"]), + ("D12", "QGD: Slav main line", START_FEN, ["d2d4", "d7d5", "c2c4", "c7c6", "g1f3", "g8f6", "b1c3", "c8f5"]), + ("D15", "QGD: Slav Gambit", START_FEN, ["d2d4", "d7d5", "c2c4", "c7c6", "b1c3", "g8f6", "g1f3", "d5c4"]), + ("D20", "QGA", START_FEN, ["d2d4", "d7d5", "c2c4", "d5c4"]), + ("D25", "QGA: Classical", START_FEN, ["d2d4", "d7d5", "c2c4", "d5c4", "g1f3", "g8f6", "e2e3", "e7e6"]), + ("D30", "QGD", START_FEN, ["d2d4", "d7d5", "c2c4", "e7e6"]), + ("D31", "QGD: Semi-Slav", START_FEN, ["d2d4", "d7d5", "c2c4", "e7e6", "b1c3", "c7c6"]), + ("D32", "QGD: Tarrasch", START_FEN, ["d2d4", "d7d5", "c2c4", "e7e6", "b1c3", "c7c5"]), + ("D35", "QGD: Exchange Variation", START_FEN, ["d2d4", "d7d5", "c2c4", "e7e6", "b1c3", "g8f6", "c4d5"]), + ("D37", "QGD: 4.Nf3", START_FEN, ["d2d4", "d7d5", "c2c4", "e7e6", "b1c3", "g8f6", "g1f3"]), + ("D41", "QGD: Semi-Tarrasch", START_FEN, ["d2d4", "d7d5", "c2c4", "e7e6", "b1c3", "g8f6", "g1f3", "c7c5", "c4d5"]), + ("D43", "QGD: Semi-Slav main", START_FEN, ["d2d4", "d7d5", "c2c4", "e7e6", "b1c3", "g8f6", "g1f3", "c7c6"]), + ("D45", "QGD: Semi-Slav Botvinnik", START_FEN, ["d2d4", "d7d5", "c2c4", "e7e6", "b1c3", "g8f6", "g1f3", "c7c6", "e2e3", "b8d7", "f1d3", "d5c4", "d3c4", "b7b5"]), + ("D50", "QGD: Semi-Slav Meran", START_FEN, ["d2d4", "d7d5", "c2c4", "e7e6", "b1c3", "g8f6", "g1f3", "c7c6", "e2e3", "b8d7", "f1d3", "d5c4", "d3c4", "b7b5", "c4d3"]), + + # ------------------------------------------------------------------ # + # King's Indian Defense (1. d4 Nf6 2. c4 g6 3. Nc3 Bg7) + # ------------------------------------------------------------------ # + ("E60", "King's Indian Defense", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6"]), + ("E61", "KID: 3.Nc3", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6", "b1c3", "f8g7"]), + ("E62", "KID: Fianchetto", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6", "b1c3", "f8g7", "g1f3", "e8g8", "g2g3"]), + ("E67", "KID: Fianchetto with ...d6", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6", "b1c3", "f8g7", "g1f3", "e8g8", "g2g3", "d7d6", "f1g2", "b8d7"]), + ("E70", "KID: Averbakh", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6", "b1c3", "f8g7", "e2e4", "d7d6"]), + ("E73", "KID: Averbakh Variation", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6", "b1c3", "f8g7", "e2e4", "d7d6", "c1e3"]), + ("E76", "KID: Four Pawns Attack", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6", "b1c3", "f8g7", "e2e4", "d7d6", "f2f4"]), + ("E80", "KID: Sämisch", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6", "b1c3", "f8g7", "e2e4", "d7d6", "f2f3"]), + ("E84", "KID: Sämisch Panno", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6", "b1c3", "f8g7", "e2e4", "d7d6", "f2f3", "b8c6", "c1e3", "a7a6"]), + ("E90", "KID: Classical", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6", "b1c3", "f8g7", "e2e4", "d7d6", "g1f3"]), + ("E92", "KID: Classical with ...e5", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6", "b1c3", "f8g7", "e2e4", "d7d6", "g1f3", "e8g8", "f1e2", "e7e5"]), + ("E97", "KID: Mar del Plata", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6", "b1c3", "f8g7", "e2e4", "d7d6", "g1f3", "e8g8", "f1e2", "e7e5", "e1g1", "b8c6"]), + + # ------------------------------------------------------------------ # + # Nimzo-Indian (1. d4 Nf6 2. c4 e6 3. Nc3 Bb4) + # ------------------------------------------------------------------ # + ("E20", "Nimzo-Indian Defense", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "b1c3", "f8b4"]), + ("E21", "Nimzo-Indian: 4.Nf3", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "b1c3", "f8b4", "g1f3"]), + ("E32", "Nimzo-Indian: Classical", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "b1c3", "f8b4", "d1c2"]), + ("E40", "Nimzo-Indian: Rubinstein", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "b1c3", "f8b4", "e2e3"]), + ("E41", "Nimzo-Indian: Huebner", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "b1c3", "f8b4", "e2e3", "c7c5"]), + ("E43", "Nimzo-Indian: Fischer", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "b1c3", "f8b4", "e2e3", "b7b6"]), + ("E46", "Nimzo-Indian: Reshevsky", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "b1c3", "f8b4", "e2e3", "e8g8", "g1e2"]), + ("E47", "Nimzo-Indian: 4.e3 d5", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "b1c3", "f8b4", "e2e3", "e8g8", "f1d3", "d7d5"]), + ("E52", "Nimzo-Indian: Spassky", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "b1c3", "f8b4", "e2e3", "e8g8", "f1d3", "d7d5", "g1f3", "c7c5"]), + ("E60b", "Nimzo-Indian: Samisch", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "b1c3", "f8b4", "a2a3"]), + + # ------------------------------------------------------------------ # + # Queen's Indian (1. d4 Nf6 2. c4 e6 3. Nf3 b6) + # ------------------------------------------------------------------ # + ("E12", "Queen's Indian Defense", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "g1f3", "b7b6"]), + ("E14", "Queen's Indian: 4.e3", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "g1f3", "b7b6", "e2e3", "c8b7"]), + ("E15", "Queen's Indian: Petrosian", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "g1f3", "b7b6", "g2g3"]), + ("E17", "Queen's Indian: Classical", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "g1f3", "b7b6", "g2g3", "c8b7", "f1g2", "f8e7"]), + + # ------------------------------------------------------------------ # + # Catalan (1. d4 Nf6 2. c4 e6 3. Nf3 d5 4. g3) + # ------------------------------------------------------------------ # + ("E00", "Catalan Opening", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "g2g3"]), + ("E01", "Catalan: Closed", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "g2g3", "d7d5", "f1g2", "f8e7"]), + ("E06", "Catalan: Open", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "g2g3", "d7d5", "g1f3", "d5c4", "f1g2"]), + + # ------------------------------------------------------------------ # + # Grünfeld Defense (1. d4 Nf6 2. c4 g6 3. Nc3 d5) + # ------------------------------------------------------------------ # + ("D80", "Grünfeld Defense", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6", "b1c3", "d7d5"]), + ("D85", "Grünfeld: Exchange", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6", "b1c3", "d7d5", "c4d5", "f6d5", "e2e4", "d5c3", "b2c3", "f8g7"]), + ("D87", "Grünfeld: Russian", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6", "b1c3", "d7d5", "c4d5", "f6d5", "e2e4", "d5c3", "b2c3", "f8g7", "f1c4"]), + ("D94", "Grünfeld: 4.e3", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6", "b1c3", "d7d5", "e2e3"]), + ("D97", "Grünfeld: Russian main line", START_FEN, ["d2d4", "g8f6", "c2c4", "g7g6", "b1c3", "d7d5", "g1f3", "f8g7", "d1b3"]), + + # ------------------------------------------------------------------ # + # English Opening (1. c4) + # ------------------------------------------------------------------ # + ("A10", "English Opening", START_FEN, ["c2c4"]), + ("A15", "English: Anglo-Indian", START_FEN, ["c2c4", "g8f6"]), + ("A17", "English: Anglo-Indian 2.Nc3", START_FEN, ["c2c4", "g8f6", "b1c3", "e7e6"]), + ("A20", "English: 1...e5", START_FEN, ["c2c4", "e7e5"]), + ("A25", "English: Closed", START_FEN, ["c2c4", "e7e5", "b1c3", "b8c6", "g2g3", "g7g6"]), + ("A29", "English: Four Knights", START_FEN, ["c2c4", "e7e5", "b1c3", "b8c6", "g1f3", "g8f6"]), + ("A30", "English: Symmetrical", START_FEN, ["c2c4", "c7c5"]), + ("A34", "English: Symmetrical 2.Nc3", START_FEN, ["c2c4", "c7c5", "b1c3", "g8f6", "g2g3"]), + + # ------------------------------------------------------------------ # + # Reti Opening (1. Nf3) + # ------------------------------------------------------------------ # + ("A04", "Reti Opening", START_FEN, ["g1f3"]), + ("A05", "Reti: 1...Nf6", START_FEN, ["g1f3", "g8f6"]), + ("A06", "Reti: 1...d5", START_FEN, ["g1f3", "d7d5"]), + ("A07", "King's Indian Attack", START_FEN, ["g1f3", "d7d5", "g2g3"]), + ("A08", "KIA with ...e5", START_FEN, ["g1f3", "d7d5", "g2g3", "g8f6", "f1g2", "c7c5"]), + + # ------------------------------------------------------------------ # + # Bird's Opening (1. f4) + # ------------------------------------------------------------------ # + ("A02", "Bird's Opening", START_FEN, ["f2f4"]), + ("A03", "Bird's: 1...d5", START_FEN, ["f2f4", "d7d5"]), + + # ------------------------------------------------------------------ # + # Dutch Defense (1. d4 f5) + # ------------------------------------------------------------------ # + ("A80", "Dutch Defense", START_FEN, ["d2d4", "f7f5"]), + ("A84", "Dutch: 2.c4", START_FEN, ["d2d4", "f7f5", "c2c4"]), + ("A87", "Dutch: Leningrad", START_FEN, ["d2d4", "f7f5", "c2c4", "g8f6", "g2g3", "g7g6"]), + ("A90", "Dutch: Classical", START_FEN, ["d2d4", "f7f5", "c2c4", "g8f6", "g2g3", "e7e6", "f1g2", "f8e7"]), + ("A92", "Dutch: Stonewall", START_FEN, ["d2d4", "f7f5", "c2c4", "g8f6", "g2g3", "e7e6", "f1g2", "f8e7", "g1f3", "e8g8", "e1g1", "d7d5"]), + + # ------------------------------------------------------------------ # + # Benoni Defense (1. d4 Nf6 2. c4 c5) + # ------------------------------------------------------------------ # + ("A56", "Benoni Defense", START_FEN, ["d2d4", "g8f6", "c2c4", "c7c5"]), + ("A60", "Modern Benoni", START_FEN, ["d2d4", "g8f6", "c2c4", "c7c5", "d4d5", "e7e6", "b1c3", "e6d5", "c4d5", "d7d6"]), + ("A65", "Benoni: 6.e4", START_FEN, ["d2d4", "g8f6", "c2c4", "c7c5", "d4d5", "e7e6", "b1c3", "e6d5", "c4d5", "d7d6", "e2e4", "g7g6"]), + ("A70", "Benoni: Classical", START_FEN, ["d2d4", "g8f6", "c2c4", "c7c5", "d4d5", "e7e6", "b1c3", "e6d5", "c4d5", "d7d6", "e2e4", "g7g6", "g1f3"]), + + # ------------------------------------------------------------------ # + # Budapest Gambit + # ------------------------------------------------------------------ # + ("A51", "Budapest Gambit", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e5"]), + + # ------------------------------------------------------------------ # + # Trompowsky Attack (1. d4 Nf6 2. Bg5) + # ------------------------------------------------------------------ # + ("A45", "Trompowsky Attack", START_FEN, ["d2d4", "g8f6", "c1g5"]), + + # ------------------------------------------------------------------ # + # Torre Attack (1. d4 Nf6 2. Nf3 e6 3. Bg5) + # ------------------------------------------------------------------ # + ("A46", "Torre Attack", START_FEN, ["d2d4", "g8f6", "g1f3", "e7e6", "c1g5"]), + + # ------------------------------------------------------------------ # + # Bogo-Indian (1. d4 Nf6 2. c4 e6 3. Nf3 Bb4+) + # ------------------------------------------------------------------ # + ("E11", "Bogo-Indian Defense", START_FEN, ["d2d4", "g8f6", "c2c4", "e7e6", "g1f3", "f8b4"]), + + # ------------------------------------------------------------------ # + # Miscellaneous + # ------------------------------------------------------------------ # + ("A00", "Grob's Attack", START_FEN, ["g2g4"]), + ("A00b", "Sokolsky Opening", START_FEN, ["b2b4"]), + ("A00c", "Van't Kruijs Opening", START_FEN, ["e2e3"]), + ("A01", "Nimzowitsch-Larsen Attack", START_FEN, ["b2b3"]), + ("A40", "Modern Defense: 1.d4 g6", START_FEN, ["d2d4", "g7g6"]), + ("A41", "Old Indian Defense", START_FEN, ["d2d4", "g7g6", "c2c4", "f8g7", "b1c3", "d7d6"]), + ("A43", "Old Benoni", START_FEN, ["d2d4", "c7c5"]), +] +# fmt: on + + +def load_openings() -> list[tuple[str, str, str, list[str]]]: + """Return the full list of (eco, name, fen, moves_uci) tuples.""" + return OPENINGS + + +def sample_opening(rng: random.Random | None = None) -> tuple[str, str, str, list[str]]: + """Return one randomly chosen opening entry.""" + r = rng or random + return r.choice(OPENINGS) diff --git a/pgn_to_training_data.py b/pgn_to_training_data.py new file mode 100644 index 0000000..a72e380 --- /dev/null +++ b/pgn_to_training_data.py @@ -0,0 +1,122 @@ +"""Parse a PGN file and produce training data for ChessFormer. + +Output format (one line per position): + <64-char board string> + +Result is from the current player's perspective: 1.0 = win, 0.5 = draw, 0.0 = loss. +Board is always oriented so the current player plays as white (flipped for black). + +For Lichess games: only games where both players have Elo >= MIN_ELO are +included, and Blitz/Bullet games are skipped. + +For self-play games (Event header starts with "SF-SelfPlay"): all games are +included regardless of Elo or time control. +""" + +import chess +import chess.pgn +import sys +from chess_moves_to_input_data import get_board_str, switch_move + +MIN_ELO = 1500 +INPUT_PGN = "lichess_db_standard_rated_2013-01.pgn" +OUTPUT_DIR = "full_datasets" +OUTPUT_FILE = f"{OUTPUT_DIR}/elo_{MIN_ELO}_pos.txt" + + +# Changed: added max_positions parameter to stop after collecting enough data +def process_pgn(pgn_path, out_path, min_elo=MIN_ELO, max_positions=None): + import os + os.makedirs(os.path.dirname(out_path), exist_ok=True) + + games_used = 0 + games_skipped = 0 + positions_written = 0 + + is_stdin = pgn_path == "/dev/stdin" + pgn_file = sys.stdin if is_stdin else open(pgn_path) + seekable = not is_stdin + + with open(out_path, "w") as out: + while True: + # Fast path: read headers first (skips move parsing), then only + # parse the full game if it passes the Elo/event filter. + if seekable: + offset = pgn_file.tell() + headers = chess.pgn.read_headers(pgn_file) + if headers is None: + break + else: + game = chess.pgn.read_game(pgn_file) + if game is None: + break + headers = game.headers + + event = headers.get("Event", "") + is_selfplay = event.startswith("SF-SelfPlay") + + if not is_selfplay: + # Filter by Elo + try: + white_elo = int(headers.get("WhiteElo", "0")) + black_elo = int(headers.get("BlackElo", "0")) + except ValueError: + games_skipped += 1 + continue + + if white_elo < min_elo or black_elo < min_elo: + games_skipped += 1 + if games_skipped % 50000 == 0: + total = games_used + games_skipped + print(f"Scanned {total} games, used: {games_used}, positions: {positions_written}") + continue + + # Skip Blitz/Bullet + if "Blitz" in event or "Bullet" in event: + games_skipped += 1 + continue + + # Game passes filter — parse moves + if seekable: + pgn_file.seek(offset) + game = chess.pgn.read_game(pgn_file) + + # Parse game result: 1-0 → 1.0 (white win), 0-1 → 0.0, 1/2-1/2 → 0.5 + result_str = game.headers.get("Result", "*") + result_map = {"1-0": 1.0, "0-1": 0.0, "1/2-1/2": 0.5} + base_result = result_map.get(result_str, 0.5) + + board = game.board() + for move in game.mainline_moves(): + board_str = get_board_str(board, white_side=board.turn) + uci = move.uci() + uci_adjusted = switch_move(uci, wht_turn=board.turn, normal_format=True) + # Result from current player's perspective (flip for black) + result = base_result if board.turn else 1.0 - base_result + out.write(f"{board_str} {uci_adjusted} {result}\n") + positions_written += 1 + board.push(move) + + games_used += 1 + if games_used % 1000 == 0: + total = games_used + games_skipped + print(f"Games used: {games_used} / {total} scanned, positions: {positions_written}") + + # Changed: stop early if we've collected enough positions + if max_positions and positions_written >= max_positions: + print(f"\nReached {max_positions:,} positions limit, stopping.") + break + + if pgn_file is not sys.stdin: + pgn_file.close() + print(f"\nDone! Games used: {games_used} (skipped: {games_skipped}), positions: {positions_written}") + print(f"Output: {out_path}") + + +if __name__ == "__main__": + pgn = sys.argv[1] if len(sys.argv) > 1 else INPUT_PGN + out = sys.argv[2] if len(sys.argv) > 2 else OUTPUT_FILE + elo = int(sys.argv[3]) if len(sys.argv) > 3 else MIN_ELO + # Changed: optional 4th arg for max positions (e.g. 10000000 for 10M) + max_pos = int(sys.argv[4]) if len(sys.argv) > 4 else None + process_pgn(pgn, out, min_elo=elo, max_positions=max_pos) diff --git a/play_against.py b/play_against.py index 1875af7..5b6f2a9 100644 --- a/play_against.py +++ b/play_against.py @@ -1,10 +1,16 @@ import chess +import chess.engine +import glob +import math +import os +import tarfile +import time import torch from inference_test import preprocess, postprocess_valid from copy import deepcopy # Model Configuration -MODEL = "2000_elo_pos_engine_best_test_whole.pth" +MODEL = "2000_elo_pos_engine.pth" # Device setup for model if torch.backends.mps.is_available(): @@ -14,14 +20,172 @@ else: device = torch.device("cpu") -model = torch.load(f'models/{MODEL}').to(device) +model = torch.load(f'models/{MODEL}', map_location="cpu").to(device) -# Initialize the chess board and move history -board = chess.Board() +# Stockfish setup — extract tar if present +_base = os.path.dirname(os.path.abspath(__file__)) +for _tar_path in glob.glob(os.path.join(_base, "*stockfish*.tar*")): + if not os.path.isdir(os.path.join(_base, "stockfish")): + with tarfile.open(_tar_path) as _tf: + _tf.extractall(_base) + for _bin in glob.glob(os.path.join(_base, "stockfish", "stockfish*")): + if os.path.isfile(_bin): + os.chmod(_bin, 0o755) + print(f"Stockfish extracted to: {os.path.join(_base, 'stockfish')}") + break + +import shutil as _shutil +import subprocess as _subprocess + +_sf_default = _shutil.which("stockfish") or "" +if not _sf_default: + for _bin in glob.glob(os.path.join(_base, "stockfish", "stockfish*")): + if os.path.isfile(_bin) and not _bin.endswith((".tar", ".zip")): + try: + _subprocess.run([_bin, "quit"], capture_output=True, timeout=5) + _sf_default = _bin + except (OSError, _subprocess.TimeoutExpired): + pass # wrong platform binary + break + +if _sf_default: + _sf_path = input(f"Path to Stockfish binary (Enter for {os.path.basename(_sf_default)}, 'n' to skip): ").strip() + if _sf_path.lower() == 'n': + _sf_path = "" + elif not _sf_path: + _sf_path = _sf_default +else: + _sf_path = input("Path to Stockfish binary (or press Enter to skip): ").strip() + +try: + sf_engine = chess.engine.SimpleEngine.popen_uci(_sf_path) if _sf_path else None +except FileNotFoundError: + sf_engine = None + print("Stockfish not found at that path. Move rating disabled.") + + +def rate_move_sf(board_before: chess.Board, move: chess.Move, depth: int = 15): + """Rate a move using Stockfish. Returns (cp_loss, label).""" + if sf_engine is None: + return None, "" + info_before = sf_engine.analyse(board_before, chess.engine.Limit(depth=depth)) + best_cp = info_before["score"].white().score(mate_score=10000) + board_after = board_before.copy() + board_after.push(move) + info_after = sf_engine.analyse(board_after, chess.engine.Limit(depth=depth)) + actual_cp = info_after["score"].white().score(mate_score=10000) + if board_before.turn == chess.WHITE: + cp_loss = best_cp - actual_cp + else: + cp_loss = actual_cp - best_cp + cp_loss = max(0, cp_loss) + if cp_loss <= 10: + label = "Excellent (!)" + elif cp_loss <= 25: + label = "Good" + elif cp_loss <= 50: + label = "Inaccuracy (?!)" + elif cp_loss <= 100: + label = "Mistake (?)" + else: + label = "Blunder (??)" + return cp_loss, label + + +LABELS = ["Excellent (!)", "Good", "Inaccuracy (?!)", "Mistake (?)", "Blunder (??)"] + +# move_log entries: (who: "You"|"AI", cp_loss: int, label: str) +move_log = [] + + +def print_summary(): + if not move_log: + return + players = list(dict.fromkeys(w for w, _, _ in move_log)) # unique, insertion order + print("\n" + "=" * 44) + print(" GAME SUMMARY") + print("=" * 44) + for who in players: + moves = [(cp, lbl) for (w, cp, lbl) in move_log if w == who] + if not moves: + continue + counts = {lbl: 0 for lbl in LABELS} + for cp, lbl in moves: + counts[lbl] += 1 + avg_cp = sum(cp for cp, _ in moves) / len(moves) + accuracy = max(0.0, min(100.0, 100 * math.exp(-avg_cp / 150))) + print(f"\n {who} ({len(moves)} moves)") + print(f" {'Excellent (!)':18s} {counts['Excellent (!)']}") + print(f" {'Good':18s} {counts['Good']}") + print(f" {'Inaccuracy (?!)':18s} {counts['Inaccuracy (?!)']}") + print(f" {'Mistake (?)':18s} {counts['Mistake (?)']}") + print(f" {'Blunder (??)':18s} {counts['Blunder (??)']}") + print(f" {'Avg cp loss':18s} {avg_cp:.1f}") + print(f" {'Accuracy':18s} {accuracy:.1f}%") + print("=" * 44) + + +# ── shared state ──────────────────────────────────────────────────────────── +board = chess.Board() made_moves = [] +move_log = [] + + +def is_over(): + return (board.is_checkmate() or board.is_stalemate() + or board.is_insufficient_material() or board.can_claim_draw()) + + +def ai_make_move(who: str, mdl=None): + """Run the model for the current position, push the move, log the rating.""" + _mdl = mdl if mdl is not None else model + input_tensors = preprocess(board) + + def predict(rep_mv=""): + with torch.no_grad(): + output = _mdl(input_tensors) + return postprocess_valid(output, board, rep_mv=rep_mv) + + uci_move = predict() -# Function to handle player's move input -def get_player_move(board): + # Avoid 3-move repetition + tmp = deepcopy(board) + tmp.push(chess.Move.from_uci(uci_move)) + if tmp.can_claim_threefold_repetition(): + uci_move = predict(rep_mv=uci_move) + + # Prioritize checkmate + for m in board.legal_moves: + tmp = deepcopy(board) + tmp.push(m) + if tmp.is_checkmate(): + uci_move = m.uci() + break + + move = chess.Move.from_uci(uci_move) + board_before = board.copy() + board.push(move) + made_moves.append(move.uci()) + cp_loss, label = rate_move_sf(board_before, move) + if label: + move_log.append((who, cp_loss, label)) + return move, cp_loss, label + + +def sf_make_move(who: str): + """Have Stockfish play the current position, push the move, log the rating.""" + result = sf_engine.play(board, chess.engine.Limit(depth=15)) + move = result.move + board_before = board.copy() + board.push(move) + made_moves.append(move.uci()) + cp_loss, label = rate_move_sf(board_before, move) + if label: + move_log.append((who, cp_loss, label)) + return move, cp_loss, label + + +def get_player_move(): while True: move_input = input("Your move (in SAN or UCI format): ") try: @@ -32,68 +196,126 @@ def get_player_move(board): except (chess.InvalidMoveError, chess.IllegalMoveError): print("Invalid move. Please try again.") continue - if move in board.legal_moves: + board_before = board.copy() board.push(move) made_moves.append(move.uci()) + cp_loss, label = rate_move_sf(board_before, move) + if label: + move_log.append(("You", cp_loss, label)) + print(f"Your move rating: {label} (cp loss: {cp_loss})") break - return board -# Game Loop -while not (board.is_checkmate() or board.is_stalemate() or board.is_insufficient_material() or board.can_claim_draw()): - print(board) - print("\nMove history:", made_moves) - - # Determining who the AI is playing as - ai_player = input("Is the AI playing as white (w) or black (b)? ").lower() - if ai_player in ['b', 'w']: - ai_player = ai_player == 'b' +# ── mode selection ─────────────────────────────────────────────────────────── +print("\n1. Human vs AI") +print("2. AI vs AI") +if sf_engine: + print("3. Model vs Stockfish") +_valid = ("1", "2", "3") if sf_engine else ("1", "2") +while True: + mode = input("Select mode: ").strip() + if mode in _valid: break + print(f"Please enter {' or '.join(_valid)}.") - print("Invalid choice. Please enter 'w' for white or 'b' for black.") +# ── AI vs AI ───────────────────────────────────────────────────────────────── +if mode == "2": + _MODEL_WHITE = "2000_elo_pos_engine.pth" + _MODEL_BLACK = "2000_elo_pos_engine.pth" + print(f"White: {_MODEL_WHITE}") + print(f"Black: {_MODEL_BLACK}") + model_white = torch.load(f'models/{_MODEL_WHITE}', map_location="cpu").to(device) + model_black = torch.load(f'models/{_MODEL_BLACK}', map_location="cpu").to(device) -# Play as human if AI is set to play as black -if ai_player: - board = get_player_move(board) + delay_str = input("Delay between moves in seconds (default 1): ").strip() + delay = float(delay_str) if delay_str else 1.0 -count = 0 -while not (board.is_checkmate() or board.is_stalemate() or board.is_insufficient_material() or board.can_claim_draw()): - # AI's turn - input_tensors = preprocess(board) - count += 1 + count = 0 + try: + while not is_over(): + count += 1 + if board.turn == chess.WHITE: + who = f"White ({_MODEL_WHITE})" + move, cp_loss, label = ai_make_move(who, mdl=model_white) + else: + who = f"Black ({_MODEL_BLACK})" + move, cp_loss, label = ai_make_move(who, mdl=model_black) + rating_str = f" → {label} (cp loss: {cp_loss})" if label else "" + print(board) + print(f"Move {count} ({who}): {move}{rating_str}") + print("\nMove history:", made_moves) + if delay > 0 and not is_over(): + time.sleep(delay) + except (KeyboardInterrupt, EOFError): + pass + finally: + print_summary() + if sf_engine: + sf_engine.quit() - def predict_move(rep_mv=""): - with torch.no_grad(): - output = model(*input_tensors) - uci_mv = postprocess_valid(output, board, rep_mv=rep_mv) - return uci_mv - - uci_move = predict_move() - - # avoiding 3-move repetition - temp_board = deepcopy(board) - temp_board.push(chess.Move.from_uci(uci_move)) - if temp_board.can_claim_threefold_repition(): - uci_move = predict_move(rep_mv=uci_move) - - # Prioritize checkmate move if available - for move in board.legal_moves: - temp_board = deepcopy(board) - temp_board.push(move) - if temp_board.is_checkmate(): - uci_move = move.uci() +# ── Model vs Stockfish ─────────────────────────────────────────────────────── +elif mode == "3": + while True: + model_color = input("Model plays as (w)hite or (b)lack? ").lower() + if model_color in ("w", "b"): break + print("Please enter 'w' or 'b'.") + model_is_white = (model_color == "w") - # Execute AI's move - move = chess.Move.from_uci(uci_move) - board.push(move) - made_moves.append(move.uci()) + delay_str = input("Delay between moves in seconds (default 1): ").strip() + delay = float(delay_str) if delay_str else 1.0 + + count = 0 + try: + while not is_over(): + count += 1 + if (board.turn == chess.WHITE) == model_is_white: + move, cp_loss, label = ai_make_move("Model") + who = "Model" + else: + move, cp_loss, label = sf_make_move("Stockfish") + who = "Stockfish" + rating_str = f" → {label} (cp loss: {cp_loss})" if label else "" + print(board) + print(f"Move {count} ({who}): {move}{rating_str}") + print("\nMove history:", made_moves) + if delay > 0 and not is_over(): + time.sleep(delay) + except (KeyboardInterrupt, EOFError): + pass + finally: + print_summary() + if sf_engine: + sf_engine.quit() - # Display board and move history +# ── Human vs AI ────────────────────────────────────────────────────────────── +else: print(board) - print(f"Predicted move {count}: {move}") - print("\nMove history:", made_moves) + while True: + ai_color = input("Is the AI playing as white (w) or black (b)? ").lower() + if ai_color in ("w", "b"): + break + print("Please enter 'w' or 'b'.") + ai_is_black = (ai_color == "b") + + if ai_is_black: + get_player_move() - if not (board.is_checkmate() or board.is_stalemate() or board.is_insufficient_material() or board.can_claim_draw()): - board = get_player_move(board) + count = 0 + try: + while not is_over(): + count += 1 + move, cp_loss, label = ai_make_move("AI") + rating_str = f" → {label} (cp loss: {cp_loss})" if label else "" + print(board) + print(f"Predicted move {count}: {move}{rating_str}") + print("\nMove history:", made_moves) + if not is_over(): + get_player_move() + except (KeyboardInterrupt, EOFError): + pass + finally: + print_summary() + if sf_engine: + sf_engine.quit() diff --git a/play_gui.py b/play_gui.py new file mode 100644 index 0000000..3fcc64e --- /dev/null +++ b/play_gui.py @@ -0,0 +1,829 @@ +import chess +import chess.engine +import glob +import math +import os +import tarfile +import torch +import pygame +import sys +from inference_test import postprocess_valid +from chessformer import ChessTransformerV2 +from model_utils import ( + detect_device, + load_model, + preprocess_board as preprocess_v2, + preprocess_board_v1 as preprocess, +) +from policy import greedy_move_v2 +from copy import deepcopy + +# --- Constants --- +SQ_SIZE = 80 +BOARD_SIZE = SQ_SIZE * 8 +BAR_W = 64 # quality bar on the left +STATUS_H = 40 +WIN_W = BAR_W + BOARD_SIZE +WIN_SIZE = (WIN_W, BOARD_SIZE + STATUS_H) + +# Colors +LIGHT_SQ = (240, 217, 181) +DARK_SQ = (181, 136, 99) +SELECTED_LIGHT = (186, 202, 68) +SELECTED_DARK = (170, 186, 58) +LAST_MOVE_LIGHT = (205, 210, 106) +LAST_MOVE_DARK = (170, 162, 58) +CHECK_COLOR = (235, 97, 80) +STATUS_BG = (48, 48, 48) +TEXT_COLOR = (220, 220, 220) +BTN_COLOR = (70, 130, 180) +BTN_HOVER = (90, 150, 200) +BTN_TEXT = (255, 255, 255) +WHITE_PIECE_COLOR = (255, 255, 255) +BLACK_PIECE_COLOR = (0, 0, 0) +BAR_BG = (28, 28, 28) + +PIECE_UNICODE = { + 'R': '\u2656', 'N': '\u2658', 'B': '\u2657', 'Q': '\u2655', 'K': '\u2654', 'P': '\u2659', + 'r': '\u265c', 'n': '\u265e', 'b': '\u265d', 'q': '\u265b', 'k': '\u265a', 'p': '\u265f', +} + +# --- Device setup (--device flag to override auto-detection) --- +import argparse +_parser = argparse.ArgumentParser() +_parser.add_argument('--device', type=str, default='auto', choices=['auto', 'cuda', 'mps', 'cpu']) +_args = _parser.parse_args() + +device = detect_device(_args.device) + +_script_dir = os.path.dirname(os.path.abspath(__file__)) +AVAILABLE_MODELS = sorted(glob.glob(os.path.join(_script_dir, "models", "*.pth"))) + +# Stockfish setup — extract tar if present +_base = _script_dir +for _tar_path in glob.glob(os.path.join(_base, "*stockfish*.tar*")): + if not os.path.isdir(os.path.join(_base, "stockfish")): + with tarfile.open(_tar_path) as _tf: + _tf.extractall(_base) + for _bin in glob.glob(os.path.join(_base, "stockfish", "stockfish*")): + if os.path.isfile(_bin): + os.chmod(_bin, 0o755) + print(f"Stockfish extracted to: {os.path.join(_base, 'stockfish')}") + break + +# Auto-detect stockfish: system PATH first, then local binary (must be runnable) +import shutil as _shutil +import subprocess as _subprocess + +_sf_default = _shutil.which("stockfish") or "" +if not _sf_default: + for _bin in glob.glob(os.path.join(_base, "stockfish", "stockfish*")): + if os.path.isfile(_bin) and not _bin.endswith((".tar", ".zip")): + try: + _subprocess.run([_bin, "quit"], capture_output=True, timeout=5) + _sf_default = _bin + except (OSError, _subprocess.TimeoutExpired): + pass # wrong platform binary + break + +if _sf_default: + _sf_path = input(f"Path to Stockfish binary (Enter for {os.path.basename(_sf_default)}, 'n' to skip): ").strip() + if _sf_path.lower() == 'n': + _sf_path = "" + elif not _sf_path: + _sf_path = _sf_default +else: + _sf_path = input("Path to Stockfish binary (or press Enter to skip): ").strip() + +try: + _sf_engine = chess.engine.SimpleEngine.popen_uci(_sf_path) if _sf_path else None +except FileNotFoundError: + _sf_engine = None + print("Stockfish not found at that path. Move rating disabled.") + +# Changed: default delay 0.1s instead of 1s +_delay_str = input("AI vs AI delay in seconds (default 0.1): ").strip() +_ai_vs_ai_delay = float(_delay_str) if _delay_str else 0.1 + +LABELS = ["Excellent (!)", "Good", "Inaccuracy (?!)", "Mistake (?)", "Blunder (??)"] + + +def print_summary(move_log): + if not move_log: + return + players = list(dict.fromkeys(w for w, _, _ in move_log)) + print("\n" + "=" * 44) + print(" GAME SUMMARY") + print("=" * 44) + for who in players: + moves = [(cp, lbl) for (w, cp, lbl) in move_log if w == who] + if not moves: + continue + counts = {lbl: 0 for lbl in LABELS} + for cp, lbl in moves: + counts[lbl] += 1 + avg_cp = sum(cp for cp, _ in moves) / len(moves) + accuracy = max(0.0, min(100.0, 100 * math.exp(-avg_cp / 150))) + print(f"\n {who} ({len(moves)} moves)") + print(f" {'Excellent (!)':18s} {counts['Excellent (!)']}") + print(f" {'Good':18s} {counts['Good']}") + print(f" {'Inaccuracy (?!)':18s} {counts['Inaccuracy (?!)']}") + print(f" {'Mistake (?)':18s} {counts['Mistake (?)']}") + print(f" {'Blunder (??)':18s} {counts['Blunder (??)']}") + print(f" {'Avg cp loss':18s} {avg_cp:.1f}") + print(f" {'Accuracy':18s} {accuracy:.1f}%") + print("=" * 44) + + +def _quality_color(q: float): + """Map q in [0,1] to RGB: 0=red, 0.5=yellow, 1=green.""" + r = int(220 * (1.0 - q)) + g = int(200 * q) + return (r, g, 30) + + +def _load_model(path): + print(f"\n=== Loading Model ===") + print(f"Path: {path}") + print(f"Device: {device}") + + m, version, _cfg = load_model(path, device) + + num_params = sum(p.numel() for p in m.parameters()) + print(f"Parameters: {num_params:,}") + print(f"Version: {version}") + + first_weight = next(m.parameters()) + print(f"First weight stats: mean={first_weight.mean().item():.6f}, std={first_weight.std().item():.6f}") + + m.eval() + print("Model loaded successfully!\n") + return m + + +class ChessGUI: + def __init__(self): + self.sf_engine = _sf_engine + self.model_white = None + self.model_black = None + pygame.init() + self.screen = pygame.display.set_mode(WIN_SIZE) + pygame.display.set_caption("Chessformer") + self.clock = pygame.time.Clock() + + # Fonts + self.piece_font = self._init_piece_font(56) + self.label_font = pygame.font.SysFont("sans", 13) + self.status_font = pygame.font.SysFont("sans", 20) + self.btn_font = pygame.font.SysFont("sans", 24, bold=True) + self.title_font = pygame.font.SysFont("sans", 48, bold=True) + self.summary_font = pygame.font.SysFont("sans", 17) + + # Game state + self.board = chess.Board() + self.made_moves = [] + self.selected_sq = None + self.legal_dests = set() + self.last_move = None + self.flipped = False + self.ai_is_black = True + self.ai_vs_ai = False + self.vs_stockfish = False + self.sf_color_screen = False + self.model_screen = True # shown first + self.model_is_white = True + self.ai_vs_ai_delay = _ai_vs_ai_delay + self.next_ai_move_at = 0 # pygame ticks for next AI move + self.game_started = False + self.game_over = False + self.status_text = "Select a model" + self.summary_done = False + + # Move quality history: list of (q in [0,1], is_white_move) + self.move_quality = [] + # Stockfish move log: list of (who, cp_loss, label) + self.move_log = [] + + # Model selection buttons — one per .pth file, stacked vertically + cx = WIN_W // 2 + mbw, mbh, mgap = 380, 48, 12 + model_start_y = 220 + self.model_btns = [] + for i, path in enumerate(AVAILABLE_MODELS): + name = os.path.splitext(os.path.basename(path))[0] + y = model_start_y + i * (mbh + mgap) + self.model_btns.append((pygame.Rect(cx - mbw // 2, y, mbw, mbh), name, path)) + + # Start screen buttons — four buttons centred in the window + bw, bh, gap = 140, 50, 10 + total = 4 * bw + 3 * gap + x0 = cx - total // 2 + self.white_btn = pygame.Rect(x0, 300, bw, bh) + self.black_btn = pygame.Rect(x0 + bw + gap, 300, bw, bh) + self.ai_vs_ai_btn = pygame.Rect(x0 + 2*(bw+gap), 300, bw, bh) + self.vs_sf_btn = pygame.Rect(x0 + 3*(bw+gap), 300, bw, bh) + # SF color sub-screen buttons + self.sf_white_btn = pygame.Rect(cx - 130, 300, 110, bh) + self.sf_black_btn = pygame.Rect(cx + 20, 300, 110, bh) + + def _init_piece_font(self, size): + # Try bundled font first (works on all platforms) + bundled = os.path.join(_script_dir, "fonts", "NotoSansSymbols2-Regular.ttf") + if os.path.isfile(bundled): + try: + return pygame.font.Font(bundled, size) + except Exception: + pass + + # Fallback: search system fonts + candidates = [ + "DejaVu Sans", "Noto Sans Symbols2", "Noto Sans Symbols", + "Symbola", "FreeSerif", "Segoe UI Symbol", "Arial Unicode MS", + "Apple Symbols", + ] + test_char = "\u2654" # ♔ White King + for name in candidates: + font = pygame.font.SysFont(name, size) + if font.get_height() > 0: + surf = font.render(test_char, True, (255, 255, 255)) + if surf.get_width() > size // 4: + return font + return pygame.font.SysFont(None, size) + + # --- Coordinate helpers --- + + def rc_to_square(self, row, col): + if self.flipped: + return chess.square(7 - col, row) + return chess.square(col, 7 - row) + + def square_to_rc(self, sq): + f, r = chess.square_file(sq), chess.square_rank(sq) + if self.flipped: + return r, 7 - f + return 7 - r, f + + def _board_x(self, col): + """Pixel x of column, accounting for the left bar.""" + return BAR_W + col * SQ_SIZE + + # --- Move quality --- + + def _eval_move(self, board_before: chess.Board, move: chess.Move) -> float: + """Rate a move using Stockfish. Returns q in [0, 1] (1=excellent, 0=blunder).""" + if self.sf_engine is None: + return 0.5 + info_before = self.sf_engine.analyse(board_before, chess.engine.Limit(depth=12)) + best_cp = info_before["score"].white().score(mate_score=10000) + board_after = board_before.copy() + board_after.push(move) + info_after = self.sf_engine.analyse(board_after, chess.engine.Limit(depth=12)) + actual_cp = info_after["score"].white().score(mate_score=10000) + if board_before.turn == chess.WHITE: + cp_loss = best_cp - actual_cp + else: + cp_loss = actual_cp - best_cp + cp_loss = max(0, cp_loss) + return max(0.0, 1.0 - cp_loss / 200.0) + + # --- Drawing --- + + def _sq_color(self, row, col, sq): + light = (row + col) % 2 == 0 + if self.board.is_check() and sq == self.board.king(self.board.turn): + return CHECK_COLOR + if sq == self.selected_sq: + return SELECTED_LIGHT if light else SELECTED_DARK + if self.last_move and sq in (self.last_move.from_square, self.last_move.to_square): + return LAST_MOVE_LIGHT if light else LAST_MOVE_DARK + return LIGHT_SQ if light else DARK_SQ + + def _draw_piece(self, piece, rect): + sym = PIECE_UNICODE[piece.symbol()] + fg = WHITE_PIECE_COLOR if piece.color == chess.WHITE else BLACK_PIECE_COLOR + outline = BLACK_PIECE_COLOR if piece.color == chess.WHITE else WHITE_PIECE_COLOR + for dx in (-1, 0, 1): + for dy in (-1, 0, 1): + if dx or dy: + s = self.piece_font.render(sym, True, outline) + self.screen.blit(s, s.get_rect(center=(rect.centerx + dx, rect.centery + dy))) + s = self.piece_font.render(sym, True, fg) + self.screen.blit(s, s.get_rect(center=rect.center)) + + def draw_board(self): + for row in range(8): + for col in range(8): + sq = self.rc_to_square(row, col) + rect = pygame.Rect(self._board_x(col), row * SQ_SIZE, SQ_SIZE, SQ_SIZE) + + pygame.draw.rect(self.screen, self._sq_color(row, col, sq), rect) + + # Legal move dots + if sq in self.legal_dests: + overlay = pygame.Surface((SQ_SIZE, SQ_SIZE), pygame.SRCALPHA) + center = (SQ_SIZE // 2, SQ_SIZE // 2) + if self.board.piece_at(sq): + pygame.draw.circle(overlay, (0, 0, 0, 50), center, SQ_SIZE // 2 - 2, 6) + else: + pygame.draw.circle(overlay, (0, 0, 0, 50), center, SQ_SIZE // 6) + self.screen.blit(overlay, rect.topleft) + + piece = self.board.piece_at(sq) + if piece: + self._draw_piece(piece, rect) + + # Coordinate labels + for i in range(8): + file_idx = i if not self.flipped else 7 - i + label = chr(ord('a') + file_idx) + color = LIGHT_SQ if i % 2 == 0 else DARK_SQ + s = self.label_font.render(label, True, color) + self.screen.blit(s, (self._board_x(i) + SQ_SIZE - s.get_width() - 2, + BOARD_SIZE - s.get_height() - 2)) + + rank_idx = i if self.flipped else 7 - i + label = str(rank_idx + 1) + color = DARK_SQ if i % 2 == 0 else LIGHT_SQ + s = self.label_font.render(label, True, color) + self.screen.blit(s, (BAR_W + 2, i * SQ_SIZE + 2)) + + def draw_quality_bar(self): + """Left panel: per-move progress bar (0->1) with numerical value.""" + pygame.draw.rect(self.screen, BAR_BG, (0, 0, BAR_W, BOARD_SIZE)) + + hdr = self.label_font.render("quality", True, (120, 120, 120)) + self.screen.blit(hdr, hdr.get_rect(center=(BAR_W // 2, 9))) + + if not self.move_quality: + pygame.draw.line(self.screen, (60, 60, 60), + (BAR_W - 1, 0), (BAR_W - 1, BOARD_SIZE)) + return + + SEG_H = 20 + pad = 5 + inner_w = BAR_W - pad * 2 + max_vis = (BOARD_SIZE - 20) // SEG_H + visible = self.move_quality[-max_vis:] + n = len(visible) + start_y = BOARD_SIZE - n * SEG_H # newest at bottom + + for i, (q, is_white) in enumerate(visible): + y = start_y + i * SEG_H + + # Track background (empty bar) + pygame.draw.rect(self.screen, (45, 45, 45), + (pad, y + 2, inner_w, SEG_H - 4), border_radius=3) + + # Filled portion — proportional to q + fill_w = max(1, int(inner_w * q)) + pygame.draw.rect(self.screen, _quality_color(q), + (pad, y + 2, fill_w, SEG_H - 4), border_radius=3) + + # Numerical value centred over the bar + val_str = f"{q:.2f}" + s = self.label_font.render(val_str, True, (240, 240, 240)) + self.screen.blit(s, s.get_rect(center=(BAR_W // 2, y + SEG_H // 2))) + + # Divider + pygame.draw.line(self.screen, (60, 60, 60), + (BAR_W - 1, 0), (BAR_W - 1, BOARD_SIZE)) + + def draw_status(self): + bar = pygame.Rect(0, BOARD_SIZE, WIN_W, STATUS_H) + pygame.draw.rect(self.screen, STATUS_BG, bar) + s = self.status_font.render(self.status_text, True, TEXT_COLOR) + self.screen.blit(s, s.get_rect(center=bar.center)) + + def draw_model_screen(self): + self.screen.fill((40, 40, 40)) + cx = WIN_W // 2 + s = self.title_font.render("Chessformer", True, TEXT_COLOR) + self.screen.blit(s, s.get_rect(center=(cx, 120))) + s = self.status_font.render("Select a model", True, TEXT_COLOR) + self.screen.blit(s, s.get_rect(center=(cx, 180))) + mouse = pygame.mouse.get_pos() + if not self.model_btns: + s = self.status_font.render("No models found in models/", True, (200, 80, 80)) + self.screen.blit(s, s.get_rect(center=(cx, 260))) + for btn, name, _path in self.model_btns: + col = BTN_HOVER if btn.collidepoint(mouse) else BTN_COLOR + pygame.draw.rect(self.screen, col, btn, border_radius=8) + s = self.btn_font.render(name, True, BTN_TEXT) + self.screen.blit(s, s.get_rect(center=btn.center)) + self.draw_status() + + def draw_start_screen(self): + self.screen.fill((40, 40, 40)) + cx = WIN_W // 2 + s = self.title_font.render("Chessformer", True, TEXT_COLOR) + self.screen.blit(s, s.get_rect(center=(cx, 150))) + s = self.status_font.render("Choose game mode", True, TEXT_COLOR) + self.screen.blit(s, s.get_rect(center=(cx, 220))) + mouse = pygame.mouse.get_pos() + sf_ok = self.sf_engine is not None + for btn, label, enabled in [ + (self.white_btn, "Play White", True), + (self.black_btn, "Play Black", True), + (self.ai_vs_ai_btn, "AI vs AI", True), + (self.vs_sf_btn, "vs Stockfish", sf_ok), + ]: + col = (BTN_HOVER if btn.collidepoint(mouse) else BTN_COLOR) if enabled else (55, 55, 55) + pygame.draw.rect(self.screen, col, btn, border_radius=8) + tc = BTN_TEXT if enabled else (110, 110, 110) + s = self.btn_font.render(label, True, tc) + self.screen.blit(s, s.get_rect(center=btn.center)) + self.draw_status() + + # --- Game logic --- + + def update_status(self): + if self.board.is_checkmate(): + winner = "Black" if self.board.turn == chess.WHITE else "White" + self.status_text = f"Checkmate! {winner} wins." + self.game_over = True + elif self.board.is_stalemate(): + self.status_text = "Draw by stalemate." + self.game_over = True + elif self.board.is_insufficient_material(): + self.status_text = "Draw — insufficient material." + self.game_over = True + elif self.board.can_claim_draw(): + self.status_text = "Draw (50-move rule / repetition)." + self.game_over = True + elif self.board.is_check(): + turn = "White" if self.board.turn else "Black" + self.status_text = f"{turn} to move — Check!" + else: + turn = "White" if self.board.turn else "Black" + self.status_text = f"{turn} to move" + + def is_ai_turn(self): + if self.ai_vs_ai or self.vs_stockfish: + return True + if self.ai_is_black: + return self.board.turn == chess.BLACK + return self.board.turn == chess.WHITE + + def ai_move(self, who: str = "AI"): + if self.game_over: + return + # In AI vs AI, use the model assigned to that color; otherwise use white model + if self.ai_vs_ai: + mdl = self.model_white if self.board.turn == chess.WHITE else self.model_black + else: + mdl = self.model_white + is_v2 = isinstance(mdl, ChessTransformerV2) + + def predict(rep_mv=""): + with torch.no_grad(): + if is_v2: + board_t, feat_t = preprocess_v2(self.board, device) + policy_logits, promo_logits, _wdl, _ply = mdl(board_t, feat_t) + move = greedy_move_v2(self.board, policy_logits[0], promo_logits[0]) + uci = move.uci() + return None if uci == rep_mv else uci + else: + input_tensors = preprocess(self.board, device) + output = mdl(input_tensors) + return postprocess_valid(output, self.board, rep_mv=rep_mv) + + uci = predict() + if uci is None: + moves = list(self.board.legal_moves) + if moves: + uci = moves[0].uci() + else: + return + + # Avoid 3-move repetition + tmp = deepcopy(self.board) + tmp.push(chess.Move.from_uci(uci)) + if tmp.can_claim_threefold_repetition(): + alt = predict(rep_mv=uci) + if alt is not None: + uci = alt + + # Prioritize checkmate + for move in self.board.legal_moves: + tmp = deepcopy(self.board) + tmp.push(move) + if tmp.is_checkmate(): + uci = move.uci() + break + + move = chess.Move.from_uci(uci) + was_white = (self.board.turn == chess.WHITE) + quality = self._eval_move(self.board, move) + self.board.push(move) + self.move_quality.append((quality, was_white)) + self.made_moves.append(move.uci()) + self.last_move = move + if self.sf_engine is not None: + cp_loss = round((1.0 - quality) * 200) + if cp_loss <= 10: label = "Excellent (!)" + elif cp_loss <= 25: label = "Good" + elif cp_loss <= 50: label = "Inaccuracy (?!)" + elif cp_loss <= 100: label = "Mistake (?)" + else: label = "Blunder (??)" + self.move_log.append((who, cp_loss, label)) + + def handle_click(self, pos): + if self.game_over or self.is_ai_turn(): + return False + + # Ignore clicks inside the left quality bar + bx = pos[0] - BAR_W + if bx < 0: + return False + + col, row = bx // SQ_SIZE, pos[1] // SQ_SIZE + if not (0 <= row < 8 and 0 <= col < 8): + return False + sq = self.rc_to_square(row, col) + + if self.selected_sq is not None: + if sq in self.legal_dests: + move = chess.Move(self.selected_sq, sq) + if move not in self.board.legal_moves: + move = chess.Move(self.selected_sq, sq, promotion=chess.QUEEN) + if move in self.board.legal_moves: + was_white = (self.board.turn == chess.WHITE) + quality = self._eval_move(self.board, move) + self.board.push(move) + self.move_quality.append((quality, was_white)) + self.made_moves.append(move.uci()) + self.last_move = move + if self.sf_engine is not None: + cp_loss = round((1.0 - quality) * 200) + if cp_loss <= 10: label = "Excellent (!)" + elif cp_loss <= 25: label = "Good" + elif cp_loss <= 50: label = "Inaccuracy (?!)" + elif cp_loss <= 100: label = "Mistake (?)" + else: label = "Blunder (??)" + self.move_log.append(("You", cp_loss, label)) + self.selected_sq = None + self.legal_dests = set() + return True + + piece = self.board.piece_at(sq) + if piece and piece.color == self.board.turn: + self.selected_sq = sq + self.legal_dests = {m.to_square for m in self.board.legal_moves + if m.from_square == sq} + else: + self.selected_sq = None + self.legal_dests = set() + else: + piece = self.board.piece_at(sq) + if piece and piece.color == self.board.turn: + self.selected_sq = sq + self.legal_dests = {m.to_square for m in self.board.legal_moves + if m.from_square == sq} + return False + + # --- SF color selection screen --- + + def draw_sf_color_screen(self): + self.screen.fill((40, 40, 40)) + cx = WIN_W // 2 + s = self.title_font.render("Chessformer", True, TEXT_COLOR) + self.screen.blit(s, s.get_rect(center=(cx, 120))) + s = self.status_font.render("Model vs Stockfish — choose model's color", True, TEXT_COLOR) + self.screen.blit(s, s.get_rect(center=(cx, 220))) + mouse = pygame.mouse.get_pos() + for btn, label in [(self.sf_white_btn, "White"), (self.sf_black_btn, "Black")]: + color = BTN_HOVER if btn.collidepoint(mouse) else BTN_COLOR + pygame.draw.rect(self.screen, color, btn, border_radius=8) + s = self.btn_font.render(label, True, BTN_TEXT) + self.screen.blit(s, s.get_rect(center=btn.center)) + self.draw_status() + + # --- Stockfish move --- + + def stockfish_move(self, who: str = "Stockfish"): + if self.game_over or self.sf_engine is None: + return + result = self.sf_engine.play(self.board, chess.engine.Limit(depth=15)) + move = result.move + was_white = (self.board.turn == chess.WHITE) + quality = self._eval_move(self.board, move) + self.board.push(move) + self.move_quality.append((quality, was_white)) + self.made_moves.append(move.uci()) + self.last_move = move + cp_loss = round((1.0 - quality) * 200) + if cp_loss <= 10: label = "Excellent (!)" + elif cp_loss <= 25: label = "Good" + elif cp_loss <= 50: label = "Inaccuracy (?!)" + elif cp_loss <= 100: label = "Mistake (?)" + else: label = "Blunder (??)" + self.move_log.append((who, cp_loss, label)) + + # --- Summary overlay --- + + def draw_summary_screen(self): + if not self.move_log: + return + + LABEL_COLORS = { + "Excellent (!)": (80, 210, 80), + "Good": (140, 210, 80), + "Inaccuracy (?!)": (220, 200, 60), + "Mistake (?)": (220, 140, 50), + "Blunder (??)": (220, 70, 70), + } + + overlay = pygame.Surface(WIN_SIZE, pygame.SRCALPHA) + overlay.fill((10, 10, 10, 215)) + self.screen.blit(overlay, (0, 0)) + + title = self.btn_font.render("Stockfish Analysis", True, TEXT_COLOR) + self.screen.blit(title, title.get_rect(center=(WIN_W // 2, 32))) + pygame.draw.line(self.screen, (80, 80, 80), (30, 56), (WIN_W - 30, 56), 1) + + players = list(dict.fromkeys(w for w, _, _ in self.move_log)) + n = len(players) + panel_w = 290 + gap = 24 + total_w = n * panel_w + (n - 1) * gap + start_x = (WIN_W - total_w) // 2 + + for i, who in enumerate(players): + moves = [(cp, lbl) for (w, cp, lbl) in self.move_log if w == who] + if not moves: + continue + counts = {lbl: 0 for lbl in LABELS} + for cp, lbl in moves: + counts[lbl] += 1 + avg_cp = sum(cp for cp, _ in moves) / len(moves) + accuracy = max(0.0, min(100.0, 100 * math.exp(-avg_cp / 150))) + + px = start_x + i * (panel_w + gap) + y = 70 + + s = self.status_font.render(f"{who} — {len(moves)} moves", True, TEXT_COLOR) + self.screen.blit(s, s.get_rect(centerx=px + panel_w // 2, y=y)) + y += 30 + pygame.draw.line(self.screen, (60, 60, 60), (px, y), (px + panel_w, y), 1) + y += 8 + + for lbl in LABELS: + s = self.summary_font.render(lbl, True, LABEL_COLORS[lbl]) + s2 = self.summary_font.render(str(counts[lbl]), True, TEXT_COLOR) + self.screen.blit(s, (px + 8, y)) + self.screen.blit(s2, s2.get_rect(right=px + panel_w - 8, y=y)) + y += 26 + + y += 4 + pygame.draw.line(self.screen, (50, 50, 50), (px, y), (px + panel_w, y), 1) + y += 8 + + s = self.summary_font.render("Avg cp loss", True, (160, 160, 160)) + s2 = self.summary_font.render(f"{avg_cp:.1f}", True, TEXT_COLOR) + self.screen.blit(s, (px + 8, y)) + self.screen.blit(s2, s2.get_rect(right=px + panel_w - 8, y=y)) + y += 32 + + acc_color = _quality_color(accuracy / 100) + s = self.btn_font.render(f"{accuracy:.1f}%", True, acc_color) + s2 = self.label_font.render("accuracy", True, (140, 140, 140)) + self.screen.blit(s, s.get_rect(centerx=px + panel_w // 2, y=y)) + self.screen.blit(s2, s2.get_rect(centerx=px + panel_w // 2, y=y + 30)) + + # --- Main loop --- + + def run(self): + running = True + need_ai_move = False + + while running: + for event in pygame.event.get(): + if event.type == pygame.QUIT: + running = False + elif event.type == pygame.MOUSEBUTTONDOWN and event.button == 1: + if self.model_screen: + for btn, name, path in self.model_btns: + if btn.collidepoint(event.pos): + self.status_text = f"Loading {name}..." + self.draw_model_screen() + pygame.display.flip() + mdl = _load_model(path) + self.model_white = mdl + self.model_black = mdl + self.model_screen = False + self.status_text = "Choose game mode" + break + elif self.sf_color_screen: + if self.sf_white_btn.collidepoint(event.pos): + self.model_is_white = True + self.vs_stockfish = True + self.sf_color_screen = False + self.game_started = True + self.update_status() + self.next_ai_move_at = pygame.time.get_ticks() + elif self.sf_black_btn.collidepoint(event.pos): + self.model_is_white = False + self.vs_stockfish = True + self.sf_color_screen = False + self.game_started = True + self.update_status() + self.next_ai_move_at = pygame.time.get_ticks() + elif not self.game_started: + if self.white_btn.collidepoint(event.pos): + self.ai_is_black = True + self.flipped = False + self.game_started = True + self.update_status() + elif self.black_btn.collidepoint(event.pos): + self.ai_is_black = False + self.flipped = True + self.game_started = True + self.update_status() + need_ai_move = True + elif self.ai_vs_ai_btn.collidepoint(event.pos): + self.ai_vs_ai = True + self.game_started = True + self.update_status() + self.next_ai_move_at = pygame.time.get_ticks() + elif self.vs_sf_btn.collidepoint(event.pos) and self.sf_engine: + self.sf_color_screen = True + self.status_text = "Choose model's color" + else: + if self.handle_click(event.pos): + self.update_status() + if not self.game_over and self.is_ai_turn(): + need_ai_move = True + + if self.model_screen: + self.draw_model_screen() + elif self.sf_color_screen: + self.draw_sf_color_screen() + elif not self.game_started: + self.draw_start_screen() + else: + # Human vs AI: trigger move via flag + if need_ai_move and not self.ai_vs_ai and not self.vs_stockfish: + self.status_text = "AI is thinking..." + self.draw_quality_bar() + self.draw_board() + self.draw_status() + pygame.display.flip() + self.ai_move("AI") + self.update_status() + need_ai_move = False + + # AI vs AI: timer-driven moves + if self.ai_vs_ai and not self.game_over: + now = pygame.time.get_ticks() + if now >= self.next_ai_move_at: + who = "White" if self.board.turn == chess.WHITE else "Black" + self.status_text = f"{who} is thinking..." + self.draw_quality_bar() + self.draw_board() + self.draw_status() + pygame.display.flip() + self.ai_move(who) + self.update_status() + self.next_ai_move_at = pygame.time.get_ticks() + int(self.ai_vs_ai_delay * 1000) + + # Model vs Stockfish: timer-driven + if self.vs_stockfish and not self.game_over: + now = pygame.time.get_ticks() + if now >= self.next_ai_move_at: + model_turn = (self.board.turn == chess.WHITE) == self.model_is_white + self.draw_quality_bar() + self.draw_board() + self.draw_status() + pygame.display.flip() + if model_turn: + self.status_text = "Model is thinking..." + self.ai_move("Model") + else: + self.status_text = "Stockfish is thinking..." + self.stockfish_move("Stockfish") + self.update_status() + self.next_ai_move_at = pygame.time.get_ticks() + int(self.ai_vs_ai_delay * 1000) + + # Print terminal summary once when game ends + if self.game_over and not self.summary_done: + print_summary(self.move_log) + self.summary_done = True + + self.draw_quality_bar() + self.draw_board() + self.draw_status() + if self.game_over: + self.draw_summary_screen() + + pygame.display.flip() + self.clock.tick(30) + + if self.sf_engine: + self.sf_engine.quit() + pygame.quit() + sys.exit() + + +if __name__ == "__main__": + gui = ChessGUI() + gui.run() diff --git a/policy.py b/policy.py new file mode 100644 index 0000000..765a47c --- /dev/null +++ b/policy.py @@ -0,0 +1,240 @@ +""" +policy.py — legal-move policy distribution for ChessFormer. + +Converts raw model logits into a proper probability distribution π(m) +over legal chess moves, including promotion handling. + +Coordinate systems +------------------ +python-chess: square = rank*8 + file (rank 0 = rank-1, file 0 = a) +model (ours): square = file + (7-rank)*8 (row 0 = rank-8, always current-player POV) + +When it is Black's turn the board is flipped vertically (rank r → 7-r) so +the model always sees its own pieces at the bottom two rows. +""" + +import chess +import torch +import torch.nn.functional as F +from typing import List, Tuple + +# Promotion piece → promo_logits index +PROMO_PIECE_TO_IDX = { + chess.QUEEN: 0, + chess.ROOK: 1, + chess.BISHOP: 2, + chess.KNIGHT: 3, +} + + +def chess_sq_to_model_idx(sq: int, white_turn: bool) -> int: + """ + Convert a python-chess square index to the model's flat board index. + + Args: + sq: python-chess square (0=a1 … 63=h8) + white_turn: True if it is White's turn (no flip needed) + + Returns: + Model board index in [0, 63] + """ + file = chess.square_file(sq) # 0-7 (a=0) + rank = chess.square_rank(sq) # 0-7 (rank-1 = 0) + if not white_turn: + rank = 7 - rank # flip for Black's perspective + return file + (7 - rank) * 8 + + +def legal_move_policy( + board: chess.Board, + from_logits: torch.Tensor, + to_logits: torch.Tensor, + promo_logits: torch.Tensor, +) -> Tuple[List[chess.Move], torch.Tensor, torch.Tensor]: + """ + Compute a softmax policy distribution over all legal moves. + + Scoring + ------- + score(m) = from_logits[f] + to_logits[t] + + promo_logits[p] (only when m is a promotion) + + where f and t are model-perspective square indices for the move's + source and destination squares. + + Args: + board: python-chess Board (any position, any side to move) + from_logits: (64,) float tensor — raw source-square logits + to_logits: (64,) float tensor — raw destination-square logits + promo_logits: (4,) float tensor — raw promotion-piece logits + order: queen=0, rook=1, bishop=2, knight=3 + + Returns: + moves: list of chess.Move in the same order as probs/log_probs + probs: (N,) tensor — π(m), softmax over legal moves + log_probs: (N,) tensor — log π(m), for use in policy-gradient / RL losses + """ + white_turn = (board.turn == chess.WHITE) + legal = list(board.legal_moves) + + if not legal: + empty = torch.zeros(0, device=from_logits.device) + return [], empty, empty + + scores = [] + for move in legal: + f = chess_sq_to_model_idx(move.from_square, white_turn) + t = chess_sq_to_model_idx(move.to_square, white_turn) + score = from_logits[f] + to_logits[t] + if move.promotion is not None: + p = PROMO_PIECE_TO_IDX[move.promotion] + score = score + promo_logits[p] + scores.append(score) + + scores = torch.stack(scores) # (N,) + log_probs = F.log_softmax(scores, dim=0) # (N,) + probs = log_probs.exp() # (N,) + + return legal, probs, log_probs + + +def sample_move( + board: chess.Board, + from_logits: torch.Tensor, + to_logits: torch.Tensor, + promo_logits: torch.Tensor, +) -> Tuple[chess.Move, torch.Tensor]: + """ + Sample a move from the policy distribution and return its log-probability. + + Args: + board, from_logits, to_logits, promo_logits: same as legal_move_policy + + Returns: + move: the sampled chess.Move + log_prob: scalar tensor — log π(move), needed for REINFORCE / PPO + """ + moves, _, log_probs = legal_move_policy(board, from_logits, to_logits, promo_logits) + idx = torch.distributions.Categorical(logits=log_probs).sample() + return moves[idx.item()], log_probs[idx] + + +def rate_move( + board: chess.Board, + move: chess.Move, + from_logits: torch.Tensor, + to_logits: torch.Tensor, + promo_logits: torch.Tensor, +) -> float: + """ + Score how much the model agrees with a move that was played. + + Returns q ∈ [0, 1] = π(move) / π(best legal move). + 1.0 → model's top choice was played (excellent) + ~0 → model strongly disagreed (poor) + """ + moves, probs, _ = legal_move_policy(board, from_logits, to_logits, promo_logits) + best_prob = probs.max().item() + if best_prob <= 0: + return 0.0 + for m, p in zip(moves, probs): + if m == move: + return p.item() / best_prob + return 0.0 # move wasn't legal (shouldn't happen) + + +def greedy_move( + board: chess.Board, + from_logits: torch.Tensor, + to_logits: torch.Tensor, + promo_logits: torch.Tensor, +) -> chess.Move: + """ + Return the highest-scoring legal move (argmax of the policy). + """ + moves, probs, _ = legal_move_policy(board, from_logits, to_logits, promo_logits) + return moves[probs.argmax().item()] + + +# --- V2: Source-destination 64x64 policy --- + + +def legal_move_policy_v2( + board: chess.Board, + policy_logits: torch.Tensor, + promo_logits: torch.Tensor, +) -> Tuple[List[chess.Move], torch.Tensor, torch.Tensor]: + """Compute softmax policy from V2 model's 64x64 source-destination logits. + + score(m) = policy_logits[from_sq, to_sq] + + promo_logits[from_sq, promo_idx] (only for promotions) + + Args: + board: python-chess Board + policy_logits: (64, 64) source-destination logit matrix + promo_logits: (64, 4) per-source promotion piece logits + + Returns: + moves, probs, log_probs — same contract as legal_move_policy + """ + white_turn = (board.turn == chess.WHITE) + legal = list(board.legal_moves) + + if not legal: + empty = torch.zeros(0, device=policy_logits.device) + return [], empty, empty + + scores = [] + for move in legal: + f = chess_sq_to_model_idx(move.from_square, white_turn) + t = chess_sq_to_model_idx(move.to_square, white_turn) + score = policy_logits[f, t] + if move.promotion is not None: + p = PROMO_PIECE_TO_IDX[move.promotion] + score = score + promo_logits[f, p] + scores.append(score) + + scores = torch.stack(scores) + log_probs = F.log_softmax(scores, dim=0) + probs = log_probs.exp() + + return legal, probs, log_probs + + +def greedy_move_v2( + board: chess.Board, + policy_logits: torch.Tensor, + promo_logits: torch.Tensor, +) -> chess.Move: + """Return highest-scoring legal move from V2 policy.""" + moves, probs, _ = legal_move_policy_v2(board, policy_logits, promo_logits) + return moves[probs.argmax().item()] + + +def sample_move_v2( + board: chess.Board, + policy_logits: torch.Tensor, + promo_logits: torch.Tensor, + temperature: float = 1.0, +) -> Tuple[chess.Move, torch.Tensor]: + """Sample a move from the V2 policy with temperature control. + + Args: + board: python-chess Board + policy_logits: (64, 64) source-destination logit matrix + promo_logits: (64, 4) per-source promotion piece logits + temperature: 0 = greedy, >1 = more random + + Returns: + move: the sampled chess.Move + log_prob: scalar tensor — log π(move) + """ + moves, probs, log_probs = legal_move_policy_v2(board, policy_logits, promo_logits) + if not moves: + raise ValueError("No legal moves available") + if temperature <= 0: + idx = probs.argmax() + else: + scaled = log_probs / temperature + idx = torch.distributions.Categorical(logits=scaled).sample() + return moves[idx.item()], log_probs[idx] diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..967b623 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,23 @@ +[build-system] +requires = ["setuptools>=61"] +build-backend = "setuptools.build_meta" + +[project] +name = "chessformer" +version = "0.1.0" +requires-python = ">=3.12" +dependencies = [ + "torch>=2.0.0", + "numpy>=1.21.0", + "python-chess>=1.999", + "pygame>=2.6.1", +] + +[tool.setuptools] +py-modules = ["chessformer"] + +[dependency-groups] +dev = [ + "pytest>=9.0.2", +] + diff --git a/requirements.txt b/requirements.txt deleted file mode 100644 index 23aba92..0000000 --- a/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -torch>=2.1.2 -torchvision>=0.16.2 -torchaudio>=2.1.2 -numpy>=1.26.3 -python-chess>=1.999 diff --git a/selfplay_loop.py b/selfplay_loop.py new file mode 100644 index 0000000..485ec47 --- /dev/null +++ b/selfplay_loop.py @@ -0,0 +1,667 @@ +"""Self-play training loop for ChessTransformerV2. + +Generates games where the V2 model plays against itself, writes positions +in the standard training format, then trains on the collected data. +Repeats for multiple generations (generate -> train -> generate -> ...). + +Usage (standalone): + uv run selfplay_loop.py --model models/2500_elo_pos_engine_v2.pth \ + --generations 10 --games-per-gen 100 --epochs-per-gen 2 + +Usage (from train_model.py): + uv run train_model.py 2500 --phase 3 \ + --backbone-model models/2500_elo_pos_engine_v2.pth \ + --generations 10 --games-per-gen 100 +""" + +from __future__ import annotations + +import random +import time +import dataclasses +from dataclasses import dataclass +from pathlib import Path + +import chess +import torch +from torch import Tensor +from torch.amp import GradScaler, autocast + +from chess_loader import get_dataloader_v2 +from chess_moves_to_input_data import get_board_str, switch_move +from chessformer import ChessTransformerV2 +from model_utils import compute_loss_v2, detect_device, find_latest_model, load_model, preprocess_board +from openings import sample_opening +from policy import sample_move_v2 + + +# --- Config --- + + +@dataclass(frozen=True) +class SelfPlayConfig: + """All parameters for one self-play training run.""" + + model_path: str + output_dir: str + generations: int + games_per_gen: int + epochs_per_gen: int + batch_size: int + lr: float + temp_schedule: list[tuple[float, int]] + max_moves: int + resign_threshold: float + resign_count: int + draw_threshold: float + draw_count: int + buffer_size: int + use_diffusion: bool + mix_supervised: str | None + mix_ratio: float + eval_games: int + device: str + mcts_sims: int # 0 = disabled (raw policy), >0 = MCTS simulations per move + cpuct: float # MCTS exploration constant + dirichlet_alpha: float # Dirichlet noise parameter (lower = spikier) + dirichlet_epsilon: float # Noise mixing weight (0 = no noise, 0.25 = AlphaZero default) + + +# --- Pure helpers --- + + +def parse_temp_schedule(s: str) -> list[tuple[float, int]]: + """Parse "1.5:10,1.0:25,0.3:999" -> [(1.5, 10), (1.0, 25), (0.3, 999)].""" + result = [] + for part in s.split(","): + temp_str, ply_str = part.strip().split(":") + result.append((float(temp_str), int(ply_str))) + return result + + +def get_temperature(schedule: list[tuple[float, int]], ply: int) -> float: + """Return temperature for the given ply based on schedule.""" + for temp, until_ply in schedule: + if ply < until_ply: + return temp + return schedule[-1][0] + + +RESULT_MAP = {"1-0": 1.0, "0-1": 0.0, "1/2-1/2": 0.5} + + +def format_time(seconds: float) -> str: + """Format seconds as '1h 23m 45s', dropping zero-value leading units.""" + s = int(seconds) + h, s = divmod(s, 3600) + m, s = divmod(s, 60) + if h > 0: + return f"{h}h {m:02d}m {s:02d}s" + if m > 0: + return f"{m}m {s:02d}s" + return f"{s}s" + + +def mcts_sample_move( + visit_policy: dict[chess.Move, float], + temperature: float, + rng: random.Random, +) -> chess.Move: + """Sample a move from MCTS visit-count policy with temperature.""" + moves = list(visit_policy.keys()) + counts = [visit_policy[m] for m in moves] + + if temperature <= 0 or len(moves) == 1: + return moves[max(range(len(counts)), key=lambda i: counts[i])] + + powered = [c ** (1.0 / temperature) for c in counts] + total = sum(powered) + probs = [p / total for p in powered] + return rng.choices(moves, weights=probs, k=1)[0] + + +# --- Game generation --- + + +def generate_game( + model: ChessTransformerV2, + device: torch.device, + opening: tuple[str, str, str, list[str]], + config: SelfPlayConfig, + rng: random.Random, +) -> list[str]: + """Play one complete self-play game, return list of training lines. + + Each line: '<64-char board> ' + Result is filled in retroactively once the game ends. + """ + _eco, _name, _fen, opening_moves = opening + + board = chess.Board() + for uci in opening_moves: + move = chess.Move.from_uci(uci) + if move not in board.legal_moves: + break # stop at first illegal opening move + board.push(move) + opening_ply = board.ply() + + positions: list[tuple[str, str, bool]] = [] # (board_str, uci_move, was_white) + resign_counter = {chess.WHITE: 0, chess.BLACK: 0} + draw_counter = 0 + resigned_side: chess.Color | None = None + + # MCTS searcher (created once per game, reused across moves) + mcts_searcher: MCTS | None = None + if config.mcts_sims > 0: + from mcts import MCTS + + mcts_searcher = MCTS( + model, device, + num_simulations=config.mcts_sims, + cpuct=config.cpuct, + use_diffusion=config.use_diffusion, + dirichlet_alpha=config.dirichlet_alpha, + dirichlet_epsilon=config.dirichlet_epsilon, + ) + + model.eval() + with torch.no_grad(): + while not board.is_game_over(claim_draw=True) and board.ply() < config.max_moves: + ply_since_opening = board.ply() - opening_ply + temperature = get_temperature(config.temp_schedule, ply_since_opening) + was_white = board.turn == chess.WHITE + + # Record position before making the move + board_str = get_board_str(board, white_side=board.turn) + + if mcts_searcher is not None: + # MCTS move selection + visit_policy, wdl_tuple = mcts_searcher.search(board) + if not visit_policy: + break # no legal moves + move = mcts_sample_move(visit_policy, temperature, rng) + loss_prob = wdl_tuple[2] + draw_prob = wdl_tuple[1] + else: + # Raw policy sampling (original behavior) + board_t, feat_t = preprocess_board(board, device) + policy_logits, promo_logits, wdl, _ply = model( + board_t, feat_t, use_diffusion=config.use_diffusion + ) + try: + move, _log_prob = sample_move_v2( + board, policy_logits[0], promo_logits[0], temperature + ) + except ValueError: + break + loss_prob = wdl[0, 2].item() + draw_prob = wdl[0, 1].item() + + # UCI from model perspective (always as-white) -> adjust for actual side + uci_adjusted = switch_move( + move.uci(), wht_turn=board.turn, normal_format=True + ) + positions.append((board_str, uci_adjusted, was_white)) + + # Adjudication: resignation + if loss_prob > config.resign_threshold: + resign_counter[board.turn] += 1 + else: + resign_counter[board.turn] = 0 + + if resign_counter[board.turn] >= config.resign_count: + resigned_side = board.turn + break + + # Adjudication: draw + if draw_prob > config.draw_threshold: + draw_counter += 1 + else: + draw_counter = 0 + + if draw_counter >= config.draw_count: + break # draw by adjudication + + board.push(move) + + # Determine result + if resigned_side is not None: + result_str = "0-1" if resigned_side == chess.WHITE else "1-0" + elif draw_counter >= config.draw_count: + result_str = "1/2-1/2" + elif board.ply() >= config.max_moves: + result_str = "1/2-1/2" + else: + outcome = board.outcome(claim_draw=True) + result_str = outcome.result() if outcome else "1/2-1/2" + + base_result = RESULT_MAP[result_str] + + # Build training lines with result from current player's perspective + lines = [] + for board_str, uci_move, was_white in positions: + result = base_result if was_white else 1.0 - base_result + lines.append(f"{board_str} {uci_move} {result}") + + return lines + + +def generate_games( + model: ChessTransformerV2, + device: torch.device, + config: SelfPlayConfig, + generation: int, +) -> str: + """Generate N games and write to output file. + + Returns path to the generated data file. + """ + output_dir = Path(config.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + data_path = output_dir / f"gen_{generation:03d}.txt" + + all_lines: list[str] = [] + results = {"1-0": 0, "0-1": 0, "1/2-1/2": 0} + rng = random.Random(generation * 1000) + + start = time.time() + for i in range(config.games_per_gen): + opening = sample_opening(rng) + lines = generate_game(model, device, opening, config, rng) + all_lines.extend(lines) + + # Count result from the lines (first position is always white's perspective) + if lines: + first_result = float(lines[0].split()[-1]) + if first_result > 0.75: + results["1-0"] += 1 + elif first_result > 0.25: + results["1/2-1/2"] += 1 + else: + results["0-1"] += 1 + + if (i + 1) % 10 == 0: + elapsed = time.time() - start + print( + f" Generated {i + 1}/{config.games_per_gen} games " + f"({len(all_lines)} positions, {format_time(elapsed)})" + ) + + with open(data_path, "w") as f: + for line in all_lines: + f.write(line + "\n") + + elapsed = time.time() - start + print( + f" Done: {config.games_per_gen} games, {len(all_lines)} positions -> {data_path}" + ) + print( + f" Results: +{results['1-0']} ={results['1/2-1/2']} -{results['0-1']} " + f"({format_time(elapsed)})" + ) + return str(data_path) + + +# --- Replay buffer --- + + +def build_training_data( + config: SelfPlayConfig, + generation: int, +) -> str: + """Concatenate recent generations + optional supervised data into one training file.""" + output_dir = Path(config.output_dir) + start_gen = max(0, generation - config.buffer_size + 1) + + all_lines: list[str] = [] + for g in range(start_gen, generation + 1): + gen_path = output_dir / f"gen_{g:03d}.txt" + with open(gen_path) as f: + all_lines.extend(f.readlines()) + + selfplay_count = len(all_lines) + + # Mix supervised data + if config.mix_supervised and config.mix_ratio > 0: + ratio = min(config.mix_ratio, 0.99) + num_supervised = int(selfplay_count * ratio / (1 - ratio)) + with open(config.mix_supervised) as f: + supervised_lines = f.readlines() + mix_rng = random.Random(generation) + sampled = mix_rng.sample( + supervised_lines, min(num_supervised, len(supervised_lines)) + ) + all_lines.extend(sampled) + + # Shuffle + random.Random(generation).shuffle(all_lines) + + combined_path = output_dir / f"combined_gen_{generation:03d}.txt" + with open(combined_path, "w") as f: + f.writelines(all_lines) + + print( + f" Training data: {selfplay_count} selfplay" + + (f" + {len(all_lines) - selfplay_count} supervised" if config.mix_supervised else "") + + f" = {len(all_lines)} total" + ) + return str(combined_path) + + +# --- Training --- + + +def train_on_selfplay( + model: ChessTransformerV2, + data_path: str, + config: SelfPlayConfig, + device: torch.device, +) -> float: + """Train model on self-play data for K epochs. Returns final train loss.""" + dataloader, testloader = get_dataloader_v2( + data_path, batch_size=config.batch_size, num_workers=4, + ) + + optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr) + use_amp = device.type == "cuda" + scaler = GradScaler("cuda") if use_amp else None + + for epoch in range(1, config.epochs_per_gen + 1): + model.train() + total_loss = 0.0 + for batch_data in dataloader: + boards, features, from_sq, to_sq, _promo, wdl_target = [ + x.to(device) for x in batch_data + ] + if use_amp: + with autocast("cuda"): + loss = compute_loss_v2(model, boards, features, from_sq, to_sq, wdl_target) + optimizer.zero_grad() + scaler.scale(loss).backward() + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1) + scaler.step(optimizer) + scaler.update() + else: + loss = compute_loss_v2(model, boards, features, from_sq, to_sq, wdl_target) + optimizer.zero_grad() + loss.backward() + torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1) + optimizer.step() + total_loss += loss.item() + + avg_train = total_loss / max(len(dataloader), 1) + + # Eval + model.eval() + test_loss = 0.0 + with torch.no_grad(): + for batch_data in testloader: + boards, features, from_sq, to_sq, _promo, wdl_target = [ + x.to(device) for x in batch_data + ] + loss = compute_loss_v2(model, boards, features, from_sq, to_sq, wdl_target) + test_loss += loss.item() + + avg_test = test_loss / max(len(testloader), 1) + print( + f" Epoch {epoch}/{config.epochs_per_gen} | " + f"Train: {avg_train:.4f} | Test: {avg_test:.4f}" + ) + + return avg_train + + +# --- Evaluation --- + + +def evaluate_model( + current: ChessTransformerV2, + baseline: ChessTransformerV2, + device: torch.device, + num_games: int, +) -> tuple[int, int, int]: + """Play evaluation games between current and baseline (greedy, no temperature). + + Returns (wins, draws, losses) from current model's perspective. + """ + from policy import greedy_move_v2 + + wins, draws, losses = 0, 0, 0 + current.eval() + baseline.eval() + + for game_idx in range(num_games): + board = chess.Board() + rng = random.Random(game_idx) + opening = sample_opening(rng) + for uci in opening[3]: + board.push(chess.Move.from_uci(uci)) + + # current plays white in even games, black in odd + current_is_white = game_idx % 2 == 0 + + with torch.no_grad(): + while not board.is_game_over(claim_draw=True) and board.ply() < 200: + is_white_turn = board.turn == chess.WHITE + active_model = ( + current if is_white_turn == current_is_white else baseline + ) + + board_t, feat_t = preprocess_board(board, device) + policy, promo, _wdl, _ply = active_model(board_t, feat_t) + move = greedy_move_v2(board, policy[0], promo[0]) + board.push(move) + + outcome = board.outcome(claim_draw=True) + if outcome is None or outcome.winner is None: + draws += 1 + elif outcome.winner == chess.WHITE: + if current_is_white: + wins += 1 + else: + losses += 1 + else: + if not current_is_white: + wins += 1 + else: + losses += 1 + + return wins, draws, losses + + +# --- Main loop --- + + +def _resolve_model_path(path: str) -> str: + """Resolve 'latest' to newest model in models/, otherwise return as-is.""" + if path == "latest": + found = find_latest_model() + if found is None: + raise FileNotFoundError("No models found in models/ directory") + return found + return path + + +def selfplay_loop(config: SelfPlayConfig) -> None: + """Main self-play training loop.""" + config = dataclasses.replace(config, model_path=_resolve_model_path(config.model_path)) + device = detect_device(config.device) + + model, _version, model_cfg = load_model(config.model_path, device) + num_params = sum(p.numel() for p in model.parameters()) + + print(f"\n{'='*60}") + print(f" Self-Play Training") + print(f"{'='*60}") + print(f" Model: {config.model_path}") + print(f" Parameters: {num_params:,}") + print(f" Device: {device}") + print(f"{'─'*60}") + print(f" Generations: {config.generations}") + print(f" Games/gen: {config.games_per_gen}") + print(f" Epochs/gen: {config.epochs_per_gen}") + print(f" Batch size: {config.batch_size}") + print(f" LR: {config.lr}") + print(f" Buffer: {config.buffer_size} generations") + print(f" Temperature: {config.temp_schedule}") + if config.mcts_sims > 0: + print(f" MCTS: {config.mcts_sims} sims, cpuct={config.cpuct}, " + f"Dir(α={config.dirichlet_alpha}, ε={config.dirichlet_epsilon})") + if config.mix_supervised: + print(f" Supervised: {config.mix_ratio:.0%} from {config.mix_supervised}") + if config.eval_games > 0: + print(f" Eval games: {config.eval_games} vs baseline") + print(f"{'='*60}\n") + + # Baseline for evaluation (frozen copy of initial model) + baseline = None + if config.eval_games > 0: + baseline, _, _ = load_model(config.model_path, device) + + # Best-model gating: track best weights, revert if new model is worse + best_state = {k: v.clone() for k, v in model.state_dict().items()} + best_gen = "baseline" + accepted = 0 + + output_dir = Path(config.output_dir) + + for gen in range(config.generations): + print(f"\n{'='*60}") + print(f" Generation {gen + 1}/{config.generations}") + print(f"{'='*60}\n") + + # 1. Generate games (always with current best model) + print("[1/4] Generating self-play games...") + generate_games(model, device, config, gen) + + # 2. Build training dataset + print("\n[2/4] Building training dataset...") + combined_path = build_training_data(config, gen) + + # 3. Train + print(f"\n[3/4] Training ({config.epochs_per_gen} epochs)...") + train_on_selfplay(model, combined_path, config, device) + + # 4. Save checkpoint (always, even if rejected later) + ckpt_path = output_dir / f"model_gen_{gen:03d}.pth" + torch.save( + { + "version": "v2", + "state_dict": model.state_dict(), + "config": model_cfg, + "selfplay_generation": gen, + }, + ckpt_path, + ) + print(f"\n Checkpoint saved: {ckpt_path}") + + # 5. Evaluate and gate + if baseline is not None and config.eval_games > 0: + print(f"\n[4/4] Evaluating ({config.eval_games} games vs baseline)...") + w, d, l = evaluate_model(model, baseline, device, config.eval_games) + total = w + d + l + score = (w + 0.5 * d) / total if total > 0 else 0.5 + # Approximate Elo difference + if 0 < score < 1: + import math + elo_diff = -400 * math.log10(1 / score - 1) + else: + elo_diff = 400 if score >= 1 else -400 + print(f" Result: +{w} ={d} -{l} (approx. {elo_diff:+.0f} Elo vs baseline)") + + # Gating: accept only if wins > losses + if w > l: + best_state = {k: v.clone() for k, v in model.state_dict().items()} + best_gen = f"gen_{gen:03d}" + accepted += 1 + print(f" -> ACCEPTED (best = {best_gen}, {accepted} accepted so far)") + else: + model.load_state_dict(best_state) + print(f" -> REJECTED, reverted to {best_gen}") + + # Save best model back to models/ directory + models_dir = Path("models") + models_dir.mkdir(exist_ok=True) + best_path = Path(config.model_path) + save_name = best_path.stem + "_selfplay" + best_path.suffix + save_path = models_dir / save_name + torch.save( + {"version": "v2", "state_dict": best_state, "config": model_cfg}, + save_path, + ) + + print(f"\n{'='*60}") + print(f"Self-play training complete! ({accepted}/{config.generations} generations accepted)") + print(f" Best model: {best_gen}") + print(f" Saved to: {save_path}") + print(f"{'='*60}") + + +# --- CLI --- + + +def main() -> None: + import argparse + + parser = argparse.ArgumentParser(description="Self-play training for ChessTransformerV2") + parser.add_argument("--model", required=True, + help='Path to V2 model checkpoint, or "latest" for newest in models/') + parser.add_argument("--output-dir", default="selfplay_data", help="Output directory") + parser.add_argument("--generations", type=int, default=20) + parser.add_argument("--games-per-gen", type=int, default=200) + parser.add_argument("--epochs-per-gen", type=int, default=3) + parser.add_argument("--batch-size", type=int, default=512) + parser.add_argument("--lr", type=float, default=1e-5) + parser.add_argument( + "--temp-schedule", type=str, default="1.5:10,1.0:25,0.3:999", + help='Temperature schedule as "temp:until_ply,..." (default: 1.5:10,1.0:25,0.3:999)', + ) + parser.add_argument("--max-moves", type=int, default=200) + parser.add_argument("--buffer-size", type=int, default=5) + parser.add_argument("--mix-supervised", type=str, default=None) + parser.add_argument("--mix-ratio", type=float, default=0.5) + parser.add_argument("--eval-games", type=int, default=100) + parser.add_argument("--mcts-sims", type=int, default=0, + help="MCTS simulations per move (0=disabled, raw policy)") + parser.add_argument("--cpuct", type=float, default=1.25, + help="MCTS exploration constant") + parser.add_argument("--dirichlet-alpha", type=float, default=0.3, + help="Dirichlet noise alpha (lower=spikier, default: 0.3)") + parser.add_argument("--dirichlet-epsilon", type=float, default=0.25, + help="Dirichlet noise weight (0=off, 0.25=AlphaZero default)") + parser.add_argument( + "--device", type=str, default="auto", choices=["auto", "cuda", "mps", "cpu"], + ) + args = parser.parse_args() + + config = SelfPlayConfig( + model_path=args.model, + output_dir=args.output_dir, + generations=args.generations, + games_per_gen=args.games_per_gen, + epochs_per_gen=args.epochs_per_gen, + batch_size=args.batch_size, + lr=args.lr, + temp_schedule=parse_temp_schedule(args.temp_schedule), + max_moves=args.max_moves, + resign_threshold=0.95, + resign_count=3, + draw_threshold=0.80, + draw_count=5, + buffer_size=args.buffer_size, + use_diffusion=False, + mix_supervised=args.mix_supervised, + mix_ratio=args.mix_ratio, + eval_games=args.eval_games, + device=args.device, + mcts_sims=args.mcts_sims, + cpuct=args.cpuct, + dirichlet_alpha=args.dirichlet_alpha, + dirichlet_epsilon=args.dirichlet_epsilon, + ) + selfplay_loop(config) + + +if __name__ == "__main__": + main() diff --git a/stockfish-ubuntu-x86-64-avx2.tar b/stockfish-ubuntu-x86-64-avx2.tar new file mode 100644 index 0000000..9762b15 --- /dev/null +++ b/stockfish-ubuntu-x86-64-avx2.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:536c0c2c0cf06450df0bfb5e876ef0d3119950703a8f143627f990c7b5417964 +size 114401280 diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/test_attention.py b/tests/test_attention.py new file mode 100644 index 0000000..17fa215 --- /dev/null +++ b/tests/test_attention.py @@ -0,0 +1,138 @@ +"""Tests for attention.py — ShawRPE, SmolgenBias, ChessTransformerBlock.""" + +import torch +import pytest +from attention import ShawRelativePositionBias, SmolgenBias, ChessTransformerBlock + + +# --- ShawRelativePositionBias --- + + +class TestShawRPE: + def setup_method(self): + self.nhead = 8 + self.shaw = ShawRelativePositionBias(self.nhead) + + def test_output_shape(self): + out = self.shaw() + assert out.shape == (self.nhead, 64, 64) + + def test_symmetric_for_same_delta(self): + """Squares with same relative (rank, file) offset get the same bias.""" + out = self.shaw() + # a1→b2 and c3→d4 both have delta (+1, +1) + # a1=0, b2=9 in chess; shaw index uses rank_idx/file_idx buffers + # Just check the bias table is indexed consistently + assert out.shape[1] == 64 + assert out.shape[2] == 64 + + def test_diagonal_consistency(self): + """All (i, i+9) pairs on the board diagonal should get the same bias + (same delta_rank=+1, delta_file=+1).""" + out = self.shaw() + # In python-chess: sq 0=a1, sq 9=b2, sq 18=c3 etc. + diagonal_pairs = [(0, 9), (9, 18), (18, 27)] + biases = [out[:, i, j] for i, j in diagonal_pairs] + for b in biases[1:]: + assert torch.allclose(b, biases[0]) + + def test_gradient_flows(self): + out = self.shaw() + loss = out.sum() + loss.backward() + assert self.shaw.bias_table.grad is not None + assert self.shaw.bias_table.grad.abs().sum() > 0 + + def test_no_input_dependency(self): + """Shaw RPE has no input — calling twice gives identical output.""" + out1 = self.shaw() + out2 = self.shaw() + assert torch.equal(out1, out2) + + +# --- SmolgenBias --- + + +class TestSmolgenBias: + def setup_method(self): + self.d_model = 64 # small for tests + self.nhead = 4 + self.smolgen = SmolgenBias(self.d_model, self.nhead) + + def test_output_shape(self): + x = torch.randn(2, 64, self.d_model) + out = self.smolgen(x) + assert out.shape == (2, self.nhead, 64, 64) + + def test_content_dependent(self): + """Different inputs produce different biases.""" + x1 = torch.randn(1, 64, self.d_model) + x2 = torch.randn(1, 64, self.d_model) + out1 = self.smolgen(x1) + out2 = self.smolgen(x2) + assert not torch.allclose(out1, out2, atol=1e-5) + + def test_batch_independent(self): + """Each batch element gets its own bias.""" + x = torch.randn(3, 64, self.d_model) + out = self.smolgen(x) + # Biases for different batch elements should differ + assert not torch.allclose(out[0], out[1], atol=1e-5) + + def test_gradient_flows(self): + x = torch.randn(1, 64, self.d_model, requires_grad=True) + out = self.smolgen(x) + loss = out.sum() + loss.backward() + assert x.grad is not None + assert x.grad.abs().sum() > 0 + + +# --- ChessTransformerBlock --- + + +class TestChessTransformerBlock: + def setup_method(self): + self.d_model = 64 + self.nhead = 4 + self.d_hid = 128 + self.block = ChessTransformerBlock(self.d_model, self.nhead, self.d_hid) + + def test_output_shape(self): + x = torch.randn(2, 64, self.d_model) + out = self.block(x) + assert out.shape == x.shape + + def test_with_attn_bias(self): + x = torch.randn(2, 64, self.d_model) + bias = torch.randn(2, self.nhead, 64, 64) + out = self.block(x, attn_bias=bias) + assert out.shape == x.shape + + def test_without_attn_bias(self): + x = torch.randn(2, 64, self.d_model) + out = self.block(x, attn_bias=None) + assert out.shape == x.shape + + def test_residual_connection(self): + """Output should differ from input (block transforms) but not be wildly different + (residual keeps magnitudes in check).""" + x = torch.randn(1, 64, self.d_model) + out = self.block(x) + assert not torch.allclose(out, x, atol=1e-5) + # Residual: output norm should be same order of magnitude as input + assert out.norm() / x.norm() < 10.0 + + def test_gradient_flows(self): + x = torch.randn(1, 64, self.d_model, requires_grad=True) + bias = torch.randn(1, self.nhead, 64, 64) + out = self.block(x, attn_bias=bias) + loss = out.sum() + loss.backward() + assert x.grad is not None + assert x.grad.abs().sum() > 0 + + def test_no_nan(self): + x = torch.randn(2, 64, self.d_model) + out = self.block(x) + assert not torch.isnan(out).any() diff --git a/tests/test_benchmark.py b/tests/test_benchmark.py new file mode 100644 index 0000000..ff7b4bd --- /dev/null +++ b/tests/test_benchmark.py @@ -0,0 +1,189 @@ +"""Tests for benchmark.py — MatchResult, game logic, integration tests.""" + +import math +import shutil + +import chess +import pytest +import torch + +from benchmark import ( + MatchResult, + benchmark_vs_model, + benchmark_vs_stockfish, + play_model_vs_model_game, + play_model_vs_stockfish_game, +) +from chessformer import ChessTransformerV2 + + +@pytest.fixture +def small_model(): + """Small V2 model for fast tests.""" + return ChessTransformerV2( + d_model=64, nhead=4, d_hid=128, nlayers=2, d_policy=32, dropout=0.0, + ) + + +@pytest.fixture +def device(): + return torch.device("cpu") + + +# --- MatchResult --- + + +class TestMatchResult: + def test_total(self): + r = MatchResult(wins=5, draws=3, losses=2, label_a="A", label_b="B") + assert r.total == 10 + + def test_score_all_wins(self): + r = MatchResult(wins=10, draws=0, losses=0, label_a="A", label_b="B") + assert r.score == pytest.approx(1.0) + + def test_score_all_losses(self): + r = MatchResult(wins=0, draws=0, losses=10, label_a="A", label_b="B") + assert r.score == pytest.approx(0.0) + + def test_score_all_draws(self): + r = MatchResult(wins=0, draws=10, losses=0, label_a="A", label_b="B") + assert r.score == pytest.approx(0.5) + + def test_score_mixed(self): + r = MatchResult(wins=3, draws=4, losses=3, label_a="A", label_b="B") + assert r.score == pytest.approx(0.5) + + def test_elo_diff_even(self): + r = MatchResult(wins=5, draws=0, losses=5, label_a="A", label_b="B") + assert r.elo_diff == pytest.approx(0.0, abs=1.0) + + def test_elo_diff_positive_for_winner(self): + r = MatchResult(wins=8, draws=1, losses=1, label_a="A", label_b="B") + assert r.elo_diff > 100 + + def test_elo_diff_negative_for_loser(self): + r = MatchResult(wins=1, draws=1, losses=8, label_a="A", label_b="B") + assert r.elo_diff < -100 + + def test_summary_contains_key_info(self): + r = MatchResult(wins=5, draws=3, losses=2, label_a="A", label_b="B") + s = r.summary() + assert "A vs B" in s + assert "+5" in s + assert "=3" in s + assert "-2" in s + assert "Elo" in s + + +# --- Model vs Model game --- + + +class TestModelVsModelGame: + def test_returns_valid_result(self, small_model, device): + result = play_model_vs_model_game( + small_model, small_model, device, a_is_white=True, game_seed=42 + ) + assert result in ("1-0", "0-1", "1/2-1/2") + + def test_different_seeds_may_differ(self, small_model, device): + results = set() + for seed in range(10): + r = play_model_vs_model_game( + small_model, small_model, device, a_is_white=seed % 2 == 0, + game_seed=seed, + ) + results.add(r) + # With random model + different openings, we should get at least 1 result type + assert len(results) >= 1 + + def test_game_terminates(self, small_model, device): + """Game should always terminate (max 200 ply enforced).""" + result = play_model_vs_model_game( + small_model, small_model, device, a_is_white=True, game_seed=0 + ) + assert result in ("1-0", "0-1", "1/2-1/2") + + +# --- Benchmark vs Model --- + + +class TestBenchmarkVsModel: + def test_runs_correct_number_of_games(self, small_model, device): + result = benchmark_vs_model( + small_model, small_model, device, num_games=4, + label_a="A", label_b="B", + ) + assert result.total == 4 + assert result.label_a == "A" + assert result.label_b == "B" + + def test_score_in_valid_range(self, small_model, device): + result = benchmark_vs_model( + small_model, small_model, device, num_games=4, + ) + assert 0.0 <= result.score <= 1.0 + + def test_wins_draws_losses_sum(self, small_model, device): + result = benchmark_vs_model( + small_model, small_model, device, num_games=6, + ) + assert result.wins + result.draws + result.losses == 6 + + +# --- Stockfish integration tests --- + + +def _has_stockfish() -> bool: + """Check if Stockfish is available.""" + if shutil.which("stockfish"): + return True + from pathlib import Path + return Path("stockfish/stockfish-ubuntu-x86-64-avx2").is_file() + + +def _get_sf_path() -> str: + from pathlib import Path + p = Path("stockfish/stockfish-ubuntu-x86-64-avx2") + if p.is_file(): + return str(p) + found = shutil.which("stockfish") + if found: + return found + raise FileNotFoundError("No Stockfish") + + +@pytest.mark.skipif(not _has_stockfish(), reason="Stockfish not available") +class TestStockfishIntegration: + def test_single_game(self, small_model, device): + sf_path = _get_sf_path() + engine = chess.engine.SimpleEngine.popen_uci(sf_path) + try: + result = play_model_vs_stockfish_game( + small_model, device, engine, + sf_depth=1, model_is_white=True, game_seed=42, + ) + assert result in ("1-0", "0-1", "1/2-1/2") + finally: + engine.quit() + + def test_benchmark_vs_stockfish(self, small_model, device): + sf_path = _get_sf_path() + result = benchmark_vs_stockfish( + small_model, device, sf_path, + skill=1, num_games=2, sf_depth=1, + ) + assert result.total == 2 + assert "Stockfish" in result.label_b + + def test_single_game_as_black(self, small_model, device): + sf_path = _get_sf_path() + engine = chess.engine.SimpleEngine.popen_uci(sf_path) + try: + result = play_model_vs_stockfish_game( + small_model, device, engine, + sf_depth=1, model_is_white=False, game_seed=99, + ) + assert result in ("1-0", "0-1", "1/2-1/2") + finally: + engine.quit() diff --git a/tests/test_chessformer_v2.py b/tests/test_chessformer_v2.py new file mode 100644 index 0000000..b6dcda2 --- /dev/null +++ b/tests/test_chessformer_v2.py @@ -0,0 +1,218 @@ +"""Tests for ChessTransformerV2 — forward pass, gradients, heads, diffusion.""" + +import torch +from torch import nn +import pytest +from chessformer import ChessTransformerV2 + + +@pytest.fixture +def model(): + """Small V2 model for fast tests.""" + return ChessTransformerV2( + d_model=64, nhead=4, d_hid=128, nlayers=2, d_policy=32, dropout=0.0, + ) + + +@pytest.fixture +def batch(): + """Fake batch: (board[B,64], features[B,14]).""" + B = 3 + board = torch.randint(0, 13, (B, 64)) + features = torch.randn(B, 14) + return board, features + + +class TestForwardPass: + """V2 forward: board + features → policy, promo, wdl, ply.""" + + def test_output_shapes(self, model, batch): + """policy[B,64,64], promo[B,64,4], wdl[B,3], ply[B,1].""" + board, features = batch + policy, promo, wdl, ply = model(board, features) + B = board.shape[0] + assert policy.shape == (B, 64, 64) + assert promo.shape == (B, 64, 4) + assert wdl.shape == (B, 3) + assert ply.shape == (B, 1) + + def test_features_affect_output(self, model, batch): + """Auxiliary features should change policy (not just board alone).""" + board, features = batch + model.eval() + with torch.no_grad(): + p_with_feat, _, _, _ = model(board, features) + p_no_feat, _, _, _ = model(board, features=None) + assert not torch.allclose(p_with_feat, p_no_feat, atol=1e-5), ( + "Features should affect policy output" + ) + + def test_wdl_is_valid_distribution(self, model, batch): + """WDL sums to 1, all values in [0, 1].""" + board, features = batch + _, _, wdl, _ = model(board, features) + sums = wdl.sum(dim=-1) + assert torch.allclose(sums, torch.ones_like(sums), atol=1e-5) + assert (wdl >= 0).all() and (wdl <= 1).all() + + def test_ply_non_negative(self, model, batch): + """Ply prediction should be non-negative.""" + board, features = batch + _, _, _, ply = model(board, features) + assert (ply >= 0).all() + + def test_no_nan_in_outputs(self, model, batch): + """No NaN in any output tensor.""" + board, features = batch + policy, promo, wdl, ply = model(board, features) + for name, t in [("policy", policy), ("promo", promo), ("wdl", wdl), ("ply", ply)]: + assert not torch.isnan(t).any(), f"NaN in {name}" + + +class TestGradients: + """Gradient flow through all model parameters.""" + + def test_all_params_receive_nonzero_gradient(self, model, batch): + """Every parameter gets non-zero gradient with proper loss.""" + board, features = batch + policy, promo, wdl, ply = model(board, features) + # wdl.sum() has zero grad (softmax sums to 1), so use CE loss + target_wdl = torch.tensor([[1.0, 0.0, 0.0]] * board.shape[0]) + wdl_loss = -(target_wdl * torch.log(wdl + 1e-8)).sum() + loss = policy.sum() + promo.sum() + wdl_loss + ply.sum() + loss.backward() + no_grad = [n for n, p in model.named_parameters() if p.grad is None or p.grad.abs().sum() == 0] + assert len(no_grad) == 0, f"Params without gradient: {no_grad}" + + +class TestDeterminism: + """Eval mode should be deterministic (no dropout).""" + + def test_eval_mode_deterministic(self, model, batch): + """Same input → same output in eval mode.""" + model.eval() + board, features = batch + with torch.no_grad(): + p1, pr1, w1, pl1 = model(board, features) + p2, pr2, w2, pl2 = model(board, features) + assert torch.equal(p1, p2) + assert torch.equal(w1, w2) + + +class TestDifferentInputs: + """Model should discriminate between different positions.""" + + def test_different_boards_different_policy(self, model): + """Two different board positions → different policy logits.""" + model.eval() + b1 = torch.randint(0, 13, (1, 64)) + b2 = torch.randint(0, 13, (1, 64)) + f = torch.randn(1, 14) + with torch.no_grad(): + p1, _, _, _ = model(b1, f) + p2, _, _, _ = model(b2, f) + assert not torch.allclose(p1, p2, atol=1e-5) + + +class TestEncode: + """encode() returns backbone latent used by all output heads.""" + + def test_encode_shape(self, model, batch): + """Latent shape: [B, 64, d_model].""" + board, features = batch + latent = model.encode(board, features) + assert latent.shape == (board.shape[0], 64, model.d_model) + + def test_encode_consistent_with_forward(self, model, batch): + """encode() latent → policy_head gives same result as forward().""" + model.eval() + board, features = batch + with torch.no_grad(): + latent = model.encode(board, features) + policy_from_encode, _ = model.policy_head(latent) + policy_from_forward, _, _, _ = model(board, features) + assert torch.equal(policy_from_encode, policy_from_forward) + + def test_encode_gradient_flows_to_embedding(self, model, batch): + """Gradients from latent reach the input embedding layer.""" + board, features = batch + latent = model.encode(board, features) + latent.sum().backward() + assert model.embedding.weight.grad is not None + assert model.embedding.weight.grad.abs().sum() > 0 + + +class TestDiffusionAttachment: + """Diffusion DiT attachment to backbone.""" + + def test_attach_creates_projection_layers(self, model): + """attach_diffusion() creates latent_to_dit and dit_to_latent.""" + from diffusion_model import ChessDiT + from noise_schedule import CosineNoiseSchedule + + d_dit = 32 + dit = ChessDiT( + d_dit=d_dit, d_model=model.d_model, + nhead=4, d_hid=64, nlayers=1, T=5, + ) + ns = CosineNoiseSchedule(T=5) + model.attach_diffusion(dit, ns, d_dit) + + assert hasattr(model, "latent_to_dit") + assert hasattr(model, "dit_to_latent") + assert model.latent_to_dit.in_features == model.d_model + assert model.latent_to_dit.out_features == d_dit + + def test_no_diffusion_attached_same_output(self, model, batch): + """use_diffusion=True without attached DiT = normal forward.""" + model.eval() + board, features = batch + with torch.no_grad(): + p1, _, _, _ = model(board, features, use_diffusion=False) + p2, _, _, _ = model(board, features, use_diffusion=True) + assert torch.equal(p1, p2) + + def test_diffusion_changes_output(self, model, batch): + """With DiT attached and non-zero weights, output changes.""" + from diffusion_model import ChessDiT + from noise_schedule import CosineNoiseSchedule + + d_dit = 32 + dit = ChessDiT( + d_dit=d_dit, d_model=model.d_model, + nhead=4, d_hid=64, nlayers=1, T=3, + ) + ns = CosineNoiseSchedule(T=3) + model.attach_diffusion(dit, ns, d_dit) + + # Break zero-init so diffusion output actually affects latent + nn.init.normal_(model.dit_to_latent.weight, std=0.1) + + model.eval() + board, features = batch + with torch.no_grad(): + p_no_diff, _, _, _ = model(board, features, use_diffusion=False) + p_with_diff, _, _, _ = model(board, features, use_diffusion=True) + assert not torch.equal(p_no_diff, p_with_diff) + + +class TestDataLoader: + """ChessDatasetV2 integration: tensors have correct shapes and values.""" + + def test_chess_loader_v2_shapes_and_wdl(self): + """Dataset returns correct tensor shapes; WDL is one-hot.""" + from chess_loader import compute_features, result_to_wdl, ChessDatasetV2 + + boards = [[0] * 64, [1] * 64] + moves = [(0, 8, -1), (32, 40, 0)] + features = [compute_features("." * 64), compute_features("P" * 64)] + wdl = [result_to_wdl(1.0), result_to_wdl(0.0)] + + ds = ChessDatasetV2(boards, moves, features, wdl) + assert len(ds) == 2 + board, feat, from_sq, to_sq, promo, wdl_t = ds[0] + assert board.shape == (64,) + assert feat.shape == (14,) + assert wdl_t.shape == (3,) + # result=1.0 → win → WDL = (1, 0, 0) + assert wdl_t.tolist() == [1.0, 0.0, 0.0] diff --git a/tests/test_diffusion_model.py b/tests/test_diffusion_model.py new file mode 100644 index 0000000..af4e40f --- /dev/null +++ b/tests/test_diffusion_model.py @@ -0,0 +1,198 @@ +"""Tests for diffusion_model.py — ChessDiT, DiTBlock, TimestepEmbedding.""" + +import torch +import pytest +from diffusion_model import ChessDiT, DiTBlock, TimestepEmbedding + + +@pytest.fixture +def dit(): + """Small DiT for fast tests.""" + return ChessDiT( + d_dit=64, d_model=128, nhead=4, d_hid=128, nlayers=2, T=20, + ) + + +@pytest.fixture +def inputs(): + B = 3 + return { + "x_t": torch.randn(B, 64, 64), # noisy latent + "t": torch.randint(0, 20, (B,)), # timesteps + "backbone_latent": torch.randn(B, 64, 128), # from policy backbone + } + + +class TestChessDiT: + def test_output_shape(self, dit, inputs): + """epsilon shape matches input x_t shape.""" + eps = dit(**inputs) + assert eps.shape == inputs["x_t"].shape + + def test_zero_init_output_is_near_zero(self, dit, inputs): + """AdaLN-Zero: final_proj is zero-init, so initial eps should be ~0.""" + dit.eval() + with torch.no_grad(): + eps = dit(**inputs) + # final_proj.weight and bias are zeros → output should be exactly zero + assert eps.abs().max() < 1e-6, f"Expected near-zero output, got max={eps.abs().max():.6f}" + + def test_zero_init_weights_are_actually_zero(self, dit): + """Verify zero-init layers have exactly zero weights at construction.""" + assert dit.final_proj.weight.abs().max() == 0.0 + assert dit.final_proj.bias.abs().max() == 0.0 + # adaLN modulation in each DiTBlock should also be zero + for i, layer in enumerate(dit.layers): + w = layer.adaLN_modulation[1].weight + assert w.abs().max() == 0.0, f"DiTBlock {i} adaLN not zero-init" + + def test_conditioning_affects_output(self, dit, inputs): + """Different backbone_latent should produce different epsilon after training.""" + optimizer = torch.optim.SGD(dit.parameters(), lr=0.1) + target = torch.randn_like(inputs["x_t"]) + for _ in range(5): + optimizer.zero_grad() + loss = (dit(**inputs) - target).pow(2).sum() + loss.backward() + optimizer.step() + + dit.eval() + with torch.no_grad(): + eps1 = dit(inputs["x_t"], inputs["t"], torch.randn(3, 64, 128)) + eps2 = dit(inputs["x_t"], inputs["t"], torch.randn(3, 64, 128)) + assert not torch.allclose(eps1, eps2, atol=1e-5), ( + "Different conditioning should produce different output after training" + ) + + def test_different_timesteps_after_training(self, dit, inputs): + """After training, different timesteps produce different outputs. + + AdaLN-Zero has a multi-step gradient unblocking chain: + - Step 1: only final_proj gets gradient (everything else zero-init) + - Step 2: final_adaLN gets gradient (final_proj now non-zero) + After step 2, final_adaLN is non-zero → conditioning (incl. timestep) + affects the output through shift/scale modulation. + """ + optimizer = torch.optim.SGD(dit.parameters(), lr=0.1) + target = torch.randn_like(inputs["x_t"]) + for _ in range(5): + optimizer.zero_grad() + loss = (dit(**inputs) - target).pow(2).sum() + loss.backward() + optimizer.step() + + dit.eval() + with torch.no_grad(): + inputs1 = {**inputs, "t": torch.tensor([0, 0, 0])} + inputs2 = {**inputs, "t": torch.tensor([19, 19, 19])} + eps1 = dit(**inputs1) + eps2 = dit(**inputs2) + assert not torch.allclose(eps1, eps2, atol=1e-5) + + def test_all_params_receive_gradient_after_training(self, dit, inputs): + """After several steps, gradient flows to all layers. + + AdaLN-Zero unblocking chain needs 3 steps minimum: + - Step 1: final_proj learns (zero → non-zero) + - Step 2: final_adaLN + block adaLN_modulations learn + - Step 3: block internals (qkv, ffn) + cond_proj + time_embed learn + """ + optimizer = torch.optim.SGD(dit.parameters(), lr=0.1) + target = torch.randn_like(inputs["x_t"]) + for _ in range(5): + optimizer.zero_grad() + loss = (dit(**inputs) - target).pow(2).sum() + loss.backward() + optimizer.step() + + optimizer.zero_grad() + loss = (dit(**inputs) - target).pow(2).sum() + loss.backward() + no_grad = [ + n for n, p in dit.named_parameters() + if p.grad is None or p.grad.abs().sum() == 0 + ] + assert len(no_grad) == 0, f"Params without gradient: {no_grad}" + + def test_no_nan(self, dit, inputs): + """No NaN in output (numerical stability).""" + eps = dit(**inputs) + assert not torch.isnan(eps).any() + + +class TestDiTBlock: + def test_output_shape(self): + """Output shape matches input shape.""" + block = DiTBlock(d_dit=64, num_heads=4, d_hid=128) + x = torch.randn(2, 64, 64) + c = torch.randn(2, 64) + out = block(x, c) + assert out.shape == x.shape + + def test_zero_init_identity(self): + """With zero-init gates, block should act as identity initially.""" + block = DiTBlock(d_dit=64, num_heads=4, d_hid=128) + x = torch.randn(1, 64, 64) + c = torch.zeros(1, 64) + out = block(x, c) + assert torch.allclose(out, x, atol=1e-5) + + def test_nonzero_conditioning_still_identity(self): + """Even with nonzero conditioning, zero-init gates keep block as identity.""" + block = DiTBlock(d_dit=64, num_heads=4, d_hid=128) + x = torch.randn(1, 64, 64) + c = torch.randn(1, 64) # nonzero conditioning + out = block(x, c) + # Gates start at zero → gate1 * attn = 0, gate2 * ffn = 0 + assert torch.allclose(out, x, atol=1e-5) + + def test_trained_block_is_not_identity(self): + """After training, block should no longer be identity.""" + block = DiTBlock(d_dit=64, num_heads=4, d_hid=128) + x = torch.randn(2, 64, 64) + c = torch.randn(2, 64) + optimizer = torch.optim.SGD(block.parameters(), lr=0.1) + target = torch.randn_like(x) + for _ in range(10): + optimizer.zero_grad() + loss = (block(x, c) - target).pow(2).sum() + loss.backward() + optimizer.step() + + block.eval() + with torch.no_grad(): + out = block(x, c) + assert not torch.allclose(out, x, atol=1e-3), ( + "Block should differ from identity after training" + ) + + +class TestTimestepEmbedding: + def test_output_shape(self): + """[B] int timesteps → [B, d_dit] float embedding.""" + embed = TimestepEmbedding(T=20, d_dit=64) + t = torch.tensor([0, 10, 20]) + out = embed(t) + assert out.shape == (3, 64) + + def test_different_timesteps_produce_different_embeddings(self): + embed = TimestepEmbedding(T=20, d_dit=64) + out = embed(torch.tensor([0, 10, 20])) + assert not torch.allclose(out[0], out[1]) + assert not torch.allclose(out[1], out[2]) + assert not torch.allclose(out[0], out[2]) + + def test_same_timestep_same_embedding(self): + """Deterministic: same t → same embedding.""" + embed = TimestepEmbedding(T=20, d_dit=64) + embed.eval() + out = embed(torch.tensor([5, 5, 5])) + assert torch.equal(out[0], out[1]) + assert torch.equal(out[1], out[2]) + + def test_embedding_is_nonzero(self): + """Timestep embeddings should not be zero vectors.""" + embed = TimestepEmbedding(T=20, d_dit=64) + out = embed(torch.tensor([0, 10, 20])) + for i in range(3): + assert out[i].abs().sum() > 0, f"Embedding for t={[0,10,20][i]} is zero" diff --git a/tests/test_grok_tracker.py b/tests/test_grok_tracker.py new file mode 100644 index 0000000..1097ddc --- /dev/null +++ b/tests/test_grok_tracker.py @@ -0,0 +1,221 @@ +"""Tests for GrokTracker and Grokfast EMA gradient filter.""" + +import pytest +import torch +from torch import nn + +from grok_tracker import GrokTracker, gradfilter_ema + + +# --- GrokTracker --- + + +class TestWeightNorms: + def test_positive_for_initialized_model(self): + model = nn.Linear(10, 5) + tracker = GrokTracker(model) + norm = tracker.compute_weight_norms() + assert norm > 0 + + def test_appends_to_history(self): + model = nn.Linear(10, 5) + tracker = GrokTracker(model) + tracker.compute_weight_norms() + tracker.compute_weight_norms() + assert len(tracker.weight_norm_history) == 2 + + def test_zero_for_zero_weights(self): + model = nn.Linear(10, 5, bias=False) + nn.init.zeros_(model.weight) + tracker = GrokTracker(model) + assert tracker.compute_weight_norms() == 0.0 + + +class TestGradientNorms: + def test_positive_after_backward(self): + model = nn.Linear(10, 5) + loss = model(torch.randn(2, 10)).sum() + loss.backward() + tracker = GrokTracker(model) + assert tracker.compute_gradient_norms() > 0 + + def test_zero_before_backward(self): + model = nn.Linear(10, 5) + tracker = GrokTracker(model) + assert tracker.compute_gradient_norms() == 0.0 + + def test_appends_to_history(self): + model = nn.Linear(10, 5) + loss = model(torch.randn(2, 10)).sum() + loss.backward() + tracker = GrokTracker(model) + tracker.compute_gradient_norms() + assert len(tracker.grad_norm_history) == 1 + + +class TestEffectiveRank: + def test_full_rank_activations(self): + tracker = GrokTracker(nn.Linear(1, 1)) + act = torch.randn(32, 16) + rank = tracker.compute_effective_rank(act) + # Random matrix has high effective rank (close to min(rows, cols)) + assert rank > 5.0 + + def test_rank_one_activations(self): + tracker = GrokTracker(nn.Linear(1, 1)) + # All columns identical → rank 1 + col = torch.randn(32, 1) + act = col.expand(32, 16) + rank = tracker.compute_effective_rank(act) + assert rank < 1.5 # close to 1 + + def test_low_rank_lower_than_full(self): + tracker = GrokTracker(nn.Linear(1, 1)) + full = torch.randn(32, 16) + low = torch.randn(32, 1).expand(32, 16) + assert tracker.compute_effective_rank(low) < tracker.compute_effective_rank(full) + + def test_3d_input(self): + """Works with [B, seq, d_model] shaped activations.""" + tracker = GrokTracker(nn.Linear(1, 1)) + act = torch.randn(4, 64, 32) # batch=4, seq=64, d_model=32 + rank = tracker.compute_effective_rank(act) + assert rank > 1.0 + + +class TestGrokkingOnset: + def test_not_enough_data(self): + tracker = GrokTracker(nn.Linear(1, 1)) + tracker.weight_norm_history = [100.0] * 5 + assert not tracker.check_grokking_onset(window=5) + + def test_detected_on_decrease(self): + tracker = GrokTracker(nn.Linear(1, 1)) + # Norms were 100, then dropped to 80 → 20% decrease > 5% threshold + tracker.weight_norm_history = [100.0] * 10 + [80.0] * 10 + assert tracker.check_grokking_onset(window=10) + + def test_not_detected_when_stable(self): + tracker = GrokTracker(nn.Linear(1, 1)) + tracker.weight_norm_history = [100.0] * 20 + assert not tracker.check_grokking_onset(window=10) + + def test_not_detected_on_increase(self): + tracker = GrokTracker(nn.Linear(1, 1)) + tracker.weight_norm_history = [80.0] * 10 + [100.0] * 10 + assert not tracker.check_grokking_onset(window=10) + + +class TestLogEpoch: + def test_returns_metrics(self): + model = nn.Linear(10, 5) + tracker = GrokTracker(model) + metrics = tracker.log_epoch(1, train_loss=2.5, test_loss=3.0) + assert "weight_norm" in metrics + assert "loss_gap" in metrics + assert "grokking_onset" in metrics + + def test_loss_gap_positive_means_overfitting(self): + tracker = GrokTracker(nn.Linear(1, 1)) + metrics = tracker.log_epoch(1, train_loss=1.0, test_loss=2.0) + # test > train → overfitting → positive gap + assert metrics["loss_gap"] == pytest.approx(1.0) + + def test_tracks_history(self): + tracker = GrokTracker(nn.Linear(1, 1)) + tracker.log_epoch(1, 2.5, 3.0) + tracker.log_epoch(2, 2.0, 2.8) + assert len(tracker.train_loss_history) == 2 + assert len(tracker.test_loss_history) == 2 + + def test_writes_to_file(self, tmp_path): + log_file = tmp_path / "grok.log" + tracker = GrokTracker(nn.Linear(1, 1), log_path=str(log_file)) + tracker.log_epoch(1, 2.5, 3.0) + tracker.close() + content = log_file.read_text() + assert "Epoch 1" in content + assert "W_norm" in content + + +class TestLogBatch: + def test_returns_both_norms(self): + model = nn.Linear(10, 5) + loss = model(torch.randn(2, 10)).sum() + loss.backward() + tracker = GrokTracker(model) + metrics = tracker.log_batch(step=0) + assert "weight_norm" in metrics + assert "gradient_norm" in metrics + assert metrics["weight_norm"] > 0 + assert metrics["gradient_norm"] > 0 + + +# --- Grokfast --- + + +class TestGrokfast: + def test_first_call_initializes(self): + model = nn.Linear(10, 5) + loss = model(torch.randn(2, 10)).sum() + loss.backward() + grads = gradfilter_ema(model, None) + assert isinstance(grads, dict) + assert len(grads) > 0 + + def test_amplifies_gradients(self): + model = nn.Linear(10, 5, bias=False) + torch.manual_seed(42) + x = torch.randn(4, 10) + + # Step 1: initialize EMA + loss = model(x).sum() + loss.backward() + ema = gradfilter_ema(model, None, alpha=0.5, lamb=1.0) + + # Step 2: apply filter — gradient should be modified + model.zero_grad() + loss = model(x).sum() + loss.backward() + orig_grad = model.weight.grad.data.clone() + ema = gradfilter_ema(model, ema, alpha=0.5, lamb=1.0) + # Gradient = original + lamb * ema → different from original + assert not torch.allclose(model.weight.grad.data, orig_grad) + + def test_no_effect_with_zero_lamb(self): + model = nn.Linear(10, 5, bias=False) + x = torch.randn(4, 10) + + loss = model(x).sum() + loss.backward() + ema = gradfilter_ema(model, None, lamb=0.0) + + model.zero_grad() + loss = model(x).sum() + loss.backward() + orig_grad = model.weight.grad.data.clone() + gradfilter_ema(model, ema, lamb=0.0) + assert torch.allclose(model.weight.grad.data, orig_grad) + + def test_ema_converges(self): + """With constant gradients, EMA converges to that gradient.""" + model = nn.Linear(10, 5, bias=False) + x = torch.randn(4, 10) + ema = None + for _ in range(100): + model.zero_grad() + loss = model(x).sum() + loss.backward() + ema = gradfilter_ema(model, ema, alpha=0.9, lamb=0.0) + + # After many steps with lamb=0, EMA should equal current grad + current_grad = model.weight.grad.data + assert torch.allclose(ema["weight"], current_grad, atol=1e-3) + + def test_skips_frozen_params(self): + model = nn.Linear(10, 5) + model.bias.requires_grad = False + loss = model(torch.randn(2, 10)).sum() + loss.backward() + grads = gradfilter_ema(model, None) + assert all("bias" not in k for k in grads) diff --git a/tests/test_mcts.py b/tests/test_mcts.py new file mode 100644 index 0000000..55974d9 --- /dev/null +++ b/tests/test_mcts.py @@ -0,0 +1,86 @@ +"""Tests for mcts.py — MCTS search, node operations, integration with model.""" + +import chess +import random +import torch +import pytest + +from chessformer import ChessTransformerV2 +from mcts import MCTSNode, MCTS + + +@pytest.fixture +def small_model(): + """Small V2 model for fast tests.""" + return ChessTransformerV2( + d_model=64, nhead=4, d_hid=128, nlayers=2, d_policy=32, dropout=0.0, + ) + + +@pytest.fixture +def device(): + return torch.device("cpu") + + +class TestMCTSNode: + def test_q_value_zero_visits(self): + node = MCTSNode(board=chess.Board()) + assert node.q_value() == 0.0 + + def test_q_value_after_updates(self): + node = MCTSNode(board=chess.Board(), visit_count=4, value_sum=2.0) + assert node.q_value() == pytest.approx(0.5) + + def test_is_expanded_empty(self): + node = MCTSNode(board=chess.Board()) + assert not node.is_expanded() + + def test_is_expanded_with_children(self): + parent = MCTSNode(board=chess.Board()) + child = MCTSNode(board=chess.Board(), parent=parent) + parent.children[chess.Move.from_uci("e2e4")] = child + assert parent.is_expanded() + + +class TestMCTSSearch: + def test_returns_legal_moves(self, small_model, device): + mcts = MCTS(small_model, device, num_simulations=10) + board = chess.Board() + policy, wdl = mcts.search(board) + legal = set(board.legal_moves) + for move in policy: + assert move in legal + + def test_visit_counts_sum_to_one(self, small_model, device): + mcts = MCTS(small_model, device, num_simulations=20) + board = chess.Board() + policy, _wdl = mcts.search(board) + assert sum(policy.values()) == pytest.approx(1.0, abs=0.01) + + def test_nonempty_from_starting_position(self, small_model, device): + mcts = MCTS(small_model, device, num_simulations=5) + policy, _wdl = mcts.search(chess.Board()) + assert len(policy) > 0 + + def test_empty_for_checkmate(self, small_model, device): + """Scholar's mate position — no legal moves.""" + board = chess.Board() + for uci in ["f2f3", "e7e5", "g2g4", "d8h4"]: + board.push(chess.Move.from_uci(uci)) + assert board.is_checkmate() + mcts = MCTS(small_model, device, num_simulations=5) + policy, _wdl = mcts.search(board) + assert len(policy) == 0 + + def test_wdl_is_tuple_of_three(self, small_model, device): + mcts = MCTS(small_model, device, num_simulations=5) + _policy, wdl = mcts.search(chess.Board()) + assert len(wdl) == 3 + assert sum(wdl) == pytest.approx(1.0, abs=0.1) + + def test_more_sims_produces_valid_policy(self, small_model, device): + """More simulations still produce a valid, non-empty policy.""" + board = chess.Board() + policy, _ = MCTS(small_model, device, num_simulations=50).search(board) + assert len(policy) > 0 + assert sum(policy.values()) == pytest.approx(1.0, abs=0.01) diff --git a/tests/test_model_utils.py b/tests/test_model_utils.py new file mode 100644 index 0000000..ef43fe5 --- /dev/null +++ b/tests/test_model_utils.py @@ -0,0 +1,214 @@ +"""Tests for model_utils — device detection, preprocessing, loss computation.""" + +import chess +import torch +import pytest +from chess_loader import PIECE_TO_INDEX +from chessformer import ChessTransformerV2 +from model_utils import ( + compute_loss_v2, + detect_device, + load_model, + preprocess_board, + preprocess_board_v1, +) + + +@pytest.fixture +def small_model(): + """Small V2 model for fast tests.""" + return ChessTransformerV2( + d_model=64, nhead=4, d_hid=128, nlayers=2, d_policy=32, dropout=0.0, + ) + + +class TestDetectDevice: + def test_cpu_explicit(self): + """'cpu' → torch.device('cpu').""" + assert detect_device("cpu") == torch.device("cpu") + + def test_auto_returns_device(self): + """'auto' resolves to a valid torch.device.""" + d = detect_device("auto") + assert isinstance(d, torch.device) + + def test_invalid_device_raises(self): + """Invalid device string → RuntimeError.""" + with pytest.raises(RuntimeError): + detect_device("nonexistent_device_xyz") + + +class TestPreprocessBoard: + def test_v2_shapes_and_types(self): + """V2: board[1,64] long, features[1,14] float.""" + board = chess.Board() + board_t, feat_t = preprocess_board(board, torch.device("cpu")) + assert board_t.shape == (1, 64) + assert feat_t.shape == (1, 14) + assert board_t.dtype == torch.long + assert feat_t.dtype == torch.float32 + + def test_v2_board_values_are_valid_piece_indices(self): + """All values should be in [0, 12] (PIECE_TO_INDEX range).""" + board = chess.Board() + board_t, _ = preprocess_board(board, torch.device("cpu")) + assert (board_t >= 0).all() and (board_t <= 12).all() + + def test_v2_starting_position_has_correct_pieces(self): + """Starting position: first 8 squares = rank 8 = opponent back rank.""" + board = chess.Board() + board_t, _ = preprocess_board(board, torch.device("cpu")) + # Board display: rank 8 (top) first → rnbqkbnr (opponent pieces, lowercase) + expected_rank8 = [PIECE_TO_INDEX[c] for c in "rnbqkbnr"] + assert board_t[0, :8].tolist() == expected_rank8 + + def test_v2_material_balance_zero_at_start(self): + """Starting position has equal material → feature[0] ≈ 0.""" + board = chess.Board() + _, feat_t = preprocess_board(board, torch.device("cpu")) + assert feat_t[0, 0].item() == pytest.approx(0.0) + + def test_v2_different_for_black(self): + """After e2e4, board from black's perspective should be flipped.""" + board = chess.Board() + white_board, _ = preprocess_board(board, torch.device("cpu")) + board.push_san("e4") + black_board, _ = preprocess_board(board, torch.device("cpu")) + assert not torch.equal(white_board, black_board) + + def test_v1_shape_and_values(self): + """V1: board[1,64] long, values in [0, 12].""" + board = chess.Board() + board_t = preprocess_board_v1(board, torch.device("cpu")) + assert board_t.shape == (1, 64) + assert board_t.dtype == torch.long + assert (board_t >= 0).all() and (board_t <= 12).all() + + def test_v1_different_for_black(self): + """V1 board flips for black's turn.""" + board = chess.Board() + white_board = preprocess_board_v1(board, torch.device("cpu")) + board.push_san("e4") + black_board = preprocess_board_v1(board, torch.device("cpu")) + assert not torch.equal(white_board, black_board) + + +class TestLoadModel: + def test_v2_roundtrip(self, small_model, tmp_path): + """Save and reload V2 checkpoint — weights should match.""" + cfg = {"d_model": 64, "nhead": 4, "d_hid": 128, "nlayers": 2, "d_policy": 32} + path = tmp_path / "test_v2.pth" + torch.save( + {"version": "v2", "state_dict": small_model.state_dict(), "config": cfg}, + path, + ) + loaded, version, loaded_cfg = load_model(str(path), torch.device("cpu")) + assert version == "v2" + assert loaded_cfg == cfg + assert isinstance(loaded, ChessTransformerV2) + # Verify weights actually match + for (n1, p1), (n2, p2) in zip( + small_model.named_parameters(), loaded.named_parameters() + ): + assert torch.equal(p1.cpu(), p2.cpu()), f"Weight mismatch in {n1}" + + def test_v2_produces_same_output(self, small_model, tmp_path): + """Loaded model should produce identical output for same input.""" + cfg = {"d_model": 64, "nhead": 4, "d_hid": 128, "nlayers": 2, "d_policy": 32} + path = tmp_path / "test_v2.pth" + torch.save( + {"version": "v2", "state_dict": small_model.state_dict(), "config": cfg}, + path, + ) + loaded, _, _ = load_model(str(path), torch.device("cpu")) + + board = torch.randint(0, 13, (1, 64)) + feat = torch.randn(1, 14) + small_model.eval() + loaded.eval() + with torch.no_grad(): + p1, _, w1, _ = small_model(board, feat) + p2, _, w2, _ = loaded(board, feat) + assert torch.equal(p1, p2) + assert torch.equal(w1, w2) + + def test_invalid_checkpoint_raises(self, tmp_path): + """Non-model file → ValueError.""" + path = tmp_path / "bad.pth" + torch.save("not a model", path) + with pytest.raises(ValueError, match="Cannot load model"): + load_model(str(path), torch.device("cpu")) + + +class TestComputeLossV2: + def test_loss_is_positive_scalar(self, small_model): + """Loss is a scalar > 0 (policy CE + WDL CE).""" + B = 4 + boards = torch.randint(0, 13, (B, 64)) + features = torch.randn(B, 14) + from_sq = torch.randint(0, 64, (B,)) + to_sq = torch.randint(0, 64, (B,)) + wdl = torch.softmax(torch.randn(B, 3), dim=-1) + + loss = compute_loss_v2(small_model, boards, features, from_sq, to_sq, wdl) + assert loss.shape == () + assert loss.item() > 0 + + def test_backward_updates_parameters(self, small_model): + """Loss.backward() produces gradients (promo/ply heads excluded — unused by loss).""" + B = 4 + boards = torch.randint(0, 13, (B, 64)) + features = torch.randn(B, 14) + from_sq = torch.randint(0, 64, (B,)) + to_sq = torch.randint(0, 64, (B,)) + wdl = torch.softmax(torch.randn(B, 3), dim=-1) + + loss = compute_loss_v2(small_model, boards, features, from_sq, to_sq, wdl) + loss.backward() + + # Params that SHOULD get gradient (everything except promo/ply heads) + no_grad = [ + n for n, p in small_model.named_parameters() + if (p.grad is None or p.grad.abs().sum() == 0) + and "promo" not in n and "ply" not in n + ] + assert len(no_grad) == 0, f"Unexpected params without gradient: {no_grad}" + + # Promo/ply heads should NOT get gradient (compute_loss_v2 ignores them) + promo_ply_with_grad = [ + n for n, p in small_model.named_parameters() + if ("promo" in n or "ply" in n) + and p.grad is not None and p.grad.abs().sum() > 0 + ] + assert len(promo_ply_with_grad) == 0, ( + f"Promo/ply params shouldn't get gradient: {promo_ply_with_grad}" + ) + + def test_loss_decreases_with_training_step(self, small_model): + """One optimizer step should reduce the loss.""" + B = 8 + boards = torch.randint(0, 13, (B, 64)) + features = torch.randn(B, 14) + from_sq = torch.randint(0, 64, (B,)) + to_sq = torch.randint(0, 64, (B,)) + wdl = torch.softmax(torch.randn(B, 3), dim=-1) + + loss_before = compute_loss_v2( + small_model, boards, features, from_sq, to_sq, wdl + ).item() + + optimizer = torch.optim.SGD(small_model.parameters(), lr=0.01) + for _ in range(5): + optimizer.zero_grad() + loss = compute_loss_v2( + small_model, boards, features, from_sq, to_sq, wdl + ) + loss.backward() + optimizer.step() + + loss_after = compute_loss_v2( + small_model, boards, features, from_sq, to_sq, wdl + ).item() + assert loss_after < loss_before, ( + f"Loss should decrease: {loss_before:.4f} → {loss_after:.4f}" + ) diff --git a/tests/test_noise_schedule.py b/tests/test_noise_schedule.py new file mode 100644 index 0000000..46e0659 --- /dev/null +++ b/tests/test_noise_schedule.py @@ -0,0 +1,89 @@ +"""Tests for noise_schedule.py — CosineNoiseSchedule.""" + +import torch +import pytest +from noise_schedule import CosineNoiseSchedule + + +@pytest.fixture +def schedule(): + return CosineNoiseSchedule(T=20) + + +class TestAlphaBar: + def test_monotonically_decreasing(self, schedule): + """alpha_bar should decrease: more noise at higher timesteps.""" + diffs = schedule.alpha_bar[1:] - schedule.alpha_bar[:-1] + assert (diffs <= 0).all() + + def test_bounds(self, schedule): + """alpha_bar should be in (0, 1).""" + assert (schedule.alpha_bar > 0).all() + assert (schedule.alpha_bar < 1).all() + + def test_near_one_at_start(self, schedule): + """alpha_bar[0] should be close to 1 (almost no noise).""" + assert schedule.alpha_bar[0] > 0.99 + + def test_near_zero_at_end(self, schedule): + """alpha_bar[T] should be close to 0 (almost pure noise).""" + assert schedule.alpha_bar[-1] < 0.05 + + def test_length(self, schedule): + """alpha_bar has T+1 entries (t=0 to t=T inclusive).""" + assert len(schedule.alpha_bar) == 21 + + +class TestQSample: + def test_output_shape(self, schedule): + x_0 = torch.randn(4, 64, 256) + t = torch.randint(0, 20, (4,)) + x_t = schedule.q_sample(x_0, t) + assert x_t.shape == x_0.shape + + def test_t0_close_to_clean(self, schedule): + """At t=0, x_t should be very close to x_0 (minimal noise).""" + x_0 = torch.randn(2, 64, 256) + t = torch.zeros(2, dtype=torch.long) + x_t = schedule.q_sample(x_0, t) + assert torch.allclose(x_t, x_0, atol=0.1) + + def test_tT_close_to_noise(self, schedule): + """At t=T, x_t should have very little signal from x_0.""" + x_0 = torch.ones(2, 64, 256) * 10.0 # large signal + t = torch.full((2,), schedule.T, dtype=torch.long) + noise = torch.randn_like(x_0) + x_t = schedule.q_sample(x_0, t, noise=noise) + # Signal should be almost gone — x_t should be close to noise + signal_ratio = (x_t - noise).abs().mean() / noise.abs().mean() + assert signal_ratio < 0.3 + + def test_different_t_different_noise(self, schedule): + """Different timesteps should produce different noise levels.""" + x_0 = torch.randn(1, 64, 256) + noise = torch.randn_like(x_0) + x_t5 = schedule.q_sample(x_0, torch.tensor([5]), noise=noise) + x_t15 = schedule.q_sample(x_0, torch.tensor([15]), noise=noise) + assert not torch.allclose(x_t5, x_t15) + + def test_custom_noise(self, schedule): + """Providing noise=zeros should give x_t = sqrt(alpha_bar) * x_0.""" + x_0 = torch.randn(2, 64, 256) + t = torch.tensor([10, 10]) + noise = torch.zeros_like(x_0) + x_t = schedule.q_sample(x_0, t, noise=noise) + expected = schedule.sqrt_alpha_bar[10] * x_0 + assert torch.allclose(x_t, expected) + + +class TestDeviceTransfer: + def test_to_cpu(self, schedule): + schedule.to(torch.device("cpu")) + assert schedule.alpha_bar.device == torch.device("cpu") + + def test_q_sample_after_move(self, schedule): + schedule.to(torch.device("cpu")) + x_0 = torch.randn(2, 64, 256) + t = torch.tensor([5, 10]) + x_t = schedule.q_sample(x_0, t) + assert x_t.shape == x_0.shape diff --git a/tests/test_selfplay_loop.py b/tests/test_selfplay_loop.py new file mode 100644 index 0000000..9c4d03a --- /dev/null +++ b/tests/test_selfplay_loop.py @@ -0,0 +1,408 @@ +"""Tests for selfplay_loop.py — game generation, data format, training integration.""" + +import chess +import random +from unittest.mock import patch, MagicMock +import torch +import pytest + +from chessformer import ChessTransformerV2 +from selfplay_loop import ( + SelfPlayConfig, + preprocess_board, + generate_game, + mcts_sample_move, + parse_temp_schedule, + get_temperature, +) +from mcts import MCTS +from policy import sample_move_v2, greedy_move_v2 + + +@pytest.fixture +def small_model(): + """Small V2 model for fast tests.""" + return ChessTransformerV2( + d_model=64, nhead=4, d_hid=128, nlayers=2, d_policy=32, dropout=0.0, + ) + + +@pytest.fixture +def device(): + return torch.device("cpu") + + +@pytest.fixture +def default_config(): + return SelfPlayConfig( + model_path="test.pth", + output_dir="/tmp/test_selfplay", + generations=1, + games_per_gen=1, + epochs_per_gen=1, + batch_size=32, + lr=1e-4, + temp_schedule=[(1.0, 999)], + max_moves=50, + resign_threshold=0.99, + resign_count=100, + draw_threshold=0.99, + draw_count=100, + buffer_size=1, + use_diffusion=False, + mix_supervised=None, + mix_ratio=0.0, + eval_games=0, + device="cpu", + mcts_sims=0, + cpuct=1.25, + dirichlet_alpha=0.3, + dirichlet_epsilon=0.25, + ) + + +# --- Temp schedule --- + + +class TestTempSchedule: + def test_parse_simple(self): + result = parse_temp_schedule("1.5:10,1.0:25,0.3:999") + assert result == [(1.5, 10), (1.0, 25), (0.3, 999)] + + def test_parse_single(self): + result = parse_temp_schedule("1.0:999") + assert result == [(1.0, 999)] + + def test_get_temperature_first_range(self): + schedule = [(1.5, 10), (1.0, 25), (0.3, 999)] + assert get_temperature(schedule, 0) == 1.5 + assert get_temperature(schedule, 9) == 1.5 + + def test_get_temperature_middle_range(self): + schedule = [(1.5, 10), (1.0, 25), (0.3, 999)] + assert get_temperature(schedule, 10) == 1.0 + assert get_temperature(schedule, 24) == 1.0 + + def test_get_temperature_last_range(self): + schedule = [(1.5, 10), (1.0, 25), (0.3, 999)] + assert get_temperature(schedule, 25) == 0.3 + assert get_temperature(schedule, 100) == 0.3 + + +# --- Preprocess --- + + +class TestPreprocessBoard: + def test_output_shapes(self, device): + board = chess.Board() + board_t, feat_t = preprocess_board(board, device) + assert board_t.shape == (1, 64) + assert feat_t.shape == (1, 14) + assert board_t.dtype == torch.long + assert feat_t.dtype == torch.float32 + + def test_after_move(self, device): + board = chess.Board() + board.push_san("e4") + board_t, feat_t = preprocess_board(board, device) + assert board_t.shape == (1, 64) + assert feat_t.shape == (1, 14) + + +# --- sample_move_v2 --- + + +class TestSampleMoveV2: + def test_returns_legal_move(self, small_model, device): + board = chess.Board() + small_model.eval() + board_t, feat_t = preprocess_board(board, device) + with torch.no_grad(): + policy, promo, _, _ = small_model(board_t, feat_t) + move, log_prob = sample_move_v2(board, policy[0], promo[0], temperature=1.0) + assert move in board.legal_moves + + def test_greedy_at_zero_temp(self, small_model, device): + board = chess.Board() + small_model.eval() + board_t, feat_t = preprocess_board(board, device) + with torch.no_grad(): + policy, promo, _, _ = small_model(board_t, feat_t) + # Greedy should be deterministic + moves = [ + sample_move_v2(board, policy[0], promo[0], temperature=0.0)[0] + for _ in range(5) + ] + assert all(m == moves[0] for m in moves) + + def test_greedy_matches_greedy_move_v2(self, small_model, device): + board = chess.Board() + small_model.eval() + board_t, feat_t = preprocess_board(board, device) + with torch.no_grad(): + policy, promo, _, _ = small_model(board_t, feat_t) + greedy = greedy_move_v2(board, policy[0], promo[0]) + sampled, _ = sample_move_v2(board, policy[0], promo[0], temperature=0.0) + assert sampled == greedy + + def test_log_prob_is_scalar(self, small_model, device): + board = chess.Board() + small_model.eval() + board_t, feat_t = preprocess_board(board, device) + with torch.no_grad(): + policy, promo, _, _ = small_model(board_t, feat_t) + _, log_prob = sample_move_v2(board, policy[0], promo[0]) + assert log_prob.dim() == 0 + assert log_prob.item() <= 0.0 + + +# --- Game generation --- + + +class TestGenerateGame: + def test_produces_training_lines(self, small_model, device, default_config): + from openings import sample_opening + + opening = sample_opening(random.Random(42)) + lines = generate_game( + model=small_model, + device=device, + opening=opening, + config=default_config, + rng=random.Random(42), + ) + assert len(lines) > 0 + + def test_training_line_format(self, small_model, device, default_config): + from openings import sample_opening + + opening = sample_opening(random.Random(42)) + lines = generate_game( + model=small_model, + device=device, + opening=opening, + config=default_config, + rng=random.Random(42), + ) + for line in lines: + parts = line.split() + assert len(parts) == 3, f"Expected 3 parts, got {len(parts)}: {line}" + board_str, uci, result = parts + assert len(board_str) == 64 + assert all(c in ".PNBRQKpnbrqk" for c in board_str) + result_f = float(result) + assert result_f in (0.0, 0.5, 1.0) + + def test_max_moves_respected(self, small_model, device): + from openings import sample_opening + + config = SelfPlayConfig( + model_path="test.pth", + output_dir="/tmp/test_selfplay", + generations=1, + games_per_gen=1, + epochs_per_gen=1, + batch_size=32, + lr=1e-4, + temp_schedule=[(1.0, 999)], + max_moves=15, + resign_threshold=0.99, + resign_count=100, + draw_threshold=0.99, + draw_count=100, + buffer_size=1, + use_diffusion=False, + mix_supervised=None, + mix_ratio=0.0, + eval_games=0, + device="cpu", + mcts_sims=0, + cpuct=1.25, + dirichlet_alpha=0.3, + dirichlet_epsilon=0.25, + ) + opening = sample_opening(random.Random(42)) + lines = generate_game( + model=small_model, + device=device, + opening=opening, + config=config, + rng=random.Random(42), + ) + # max_moves=15 means max 15 total ply including opening + assert len(lines) <= 15 + + +# --- Config --- + + +class TestMctsSampleMove: + """Tests for mcts_sample_move() helper.""" + + @pytest.fixture + def sample_policy(self): + return { + chess.Move.from_uci("e2e4"): 0.6, + chess.Move.from_uci("d2d4"): 0.3, + chess.Move.from_uci("g1f3"): 0.1, + } + + def test_greedy_picks_highest(self, sample_policy): + move = mcts_sample_move(sample_policy, temperature=0.0, rng=random.Random(42)) + assert move == chess.Move.from_uci("e2e4") + + def test_greedy_deterministic(self, sample_policy): + moves = [ + mcts_sample_move(sample_policy, temperature=0.0, rng=random.Random(i)) + for i in range(10) + ] + assert all(m == moves[0] for m in moves) + + def test_sampling_returns_valid_move(self, sample_policy): + for seed in range(20): + move = mcts_sample_move(sample_policy, temperature=1.0, rng=random.Random(seed)) + assert move in sample_policy + + def test_single_move_returns_it(self): + policy = {chess.Move.from_uci("a2a3"): 1.0} + move = mcts_sample_move(policy, temperature=1.0, rng=random.Random(0)) + assert move == chess.Move.from_uci("a2a3") + + def test_high_temp_explores_more(self, sample_policy): + """High temperature should produce more variety than greedy.""" + moves_seen = set() + for seed in range(50): + move = mcts_sample_move(sample_policy, temperature=2.0, rng=random.Random(seed)) + moves_seen.add(move) + assert len(moves_seen) >= 2 + + def test_with_real_mcts_output(self, small_model, device): + """mcts_sample_move works with actual MCTS search output.""" + board = chess.Board() + mcts = MCTS(small_model, device, num_simulations=10) + visit_policy, _wdl = mcts.search(board) + move = mcts_sample_move(visit_policy, temperature=1.0, rng=random.Random(42)) + assert move in board.legal_moves + + +# --- MCTS + generate_game integration --- + + +@pytest.fixture +def mcts_config(): + return SelfPlayConfig( + model_path="test.pth", + output_dir="/tmp/test_selfplay_mcts", + generations=1, + games_per_gen=1, + epochs_per_gen=1, + batch_size=32, + lr=1e-4, + temp_schedule=[(1.0, 999)], + max_moves=30, + resign_threshold=0.99, + resign_count=100, + draw_threshold=0.99, + draw_count=100, + buffer_size=1, + use_diffusion=False, + mix_supervised=None, + mix_ratio=0.0, + eval_games=0, + device="cpu", + mcts_sims=5, + cpuct=1.25, + dirichlet_alpha=0.3, + dirichlet_epsilon=0.25, + ) + + +class TestGenerateGameMCTS: + """generate_game() with MCTS enabled — verify MCTS is actually used.""" + + def test_mcts_search_is_called(self, small_model, device, mcts_config): + """Verify MCTS.search() is actually invoked during game generation.""" + from openings import sample_opening + + opening = sample_opening(random.Random(42)) + original_search = MCTS.search + search_calls: list[chess.Board] = [] + + def tracking_search(self, board): + search_calls.append(board) + return original_search(self, board) + + with patch.object(MCTS, "search", tracking_search): + generate_game( + model=small_model, + device=device, + opening=opening, + config=mcts_config, + rng=random.Random(42), + ) + assert len(search_calls) > 0, "MCTS.search() was never called" + + def test_sample_move_v2_not_called(self, small_model, device, mcts_config): + """When MCTS is on, raw sample_move_v2 should NOT be used.""" + from openings import sample_opening + + opening = sample_opening(random.Random(42)) + with patch("selfplay_loop.sample_move_v2") as mock_sample: + generate_game( + model=small_model, + device=device, + opening=opening, + config=mcts_config, + rng=random.Random(42), + ) + mock_sample.assert_not_called() + + def test_training_line_format_with_mcts(self, small_model, device, mcts_config): + """MCTS game produces valid training lines.""" + from openings import sample_opening + + opening = sample_opening(random.Random(42)) + lines = generate_game( + model=small_model, + device=device, + opening=opening, + config=mcts_config, + rng=random.Random(42), + ) + assert len(lines) > 0 + for line in lines: + parts = line.split() + assert len(parts) == 3, f"Expected 3 parts, got {len(parts)}: {line}" + board_str, uci, result = parts + assert len(board_str) == 64 + assert all(c in ".PNBRQKpnbrqk" for c in board_str) + result_f = float(result) + assert result_f in (0.0, 0.5, 1.0) + + def test_non_mcts_uses_sample_move_v2(self, small_model, device, default_config): + """When mcts_sims=0, sample_move_v2 IS called (not MCTS).""" + from openings import sample_opening + + opening = sample_opening(random.Random(42)) + with patch("selfplay_loop.sample_move_v2", wraps=sample_move_v2) as mock_sample: + generate_game( + model=small_model, + device=device, + opening=opening, + config=default_config, + rng=random.Random(42), + ) + assert mock_sample.call_count > 0, "sample_move_v2 was never called" + + +# --- Config --- + + +class TestSelfPlayConfig: + def test_frozen(self, default_config): + with pytest.raises(AttributeError): + default_config.generations = 5 + + def test_mcts_fields_present(self, default_config): + assert default_config.mcts_sims == 0 + assert default_config.cpuct == 1.25 diff --git a/tests/test_trajectory_loader.py b/tests/test_trajectory_loader.py new file mode 100644 index 0000000..ca4335a --- /dev/null +++ b/tests/test_trajectory_loader.py @@ -0,0 +1,104 @@ +"""Tests for trajectory_loader.py — TrajectoryDataset, extract_trajectories.""" + +import torch +import pytest +from chess_loader import PIECE_TO_INDEX, compute_features +from trajectory_loader import TrajectoryDataset + + +# Starting position and after 1.e4 (from white's perspective) +STARTING_BOARD = "RNBQKBNRPPPPPPPP................................pppppppprnbqkbnr" +AFTER_E4 = "RNBQKBNRPPPP.PPP............P...................pppppppprnbqkbnr" +# After 1...e5 (from black's perspective — board is flipped) +AFTER_E4_E5 = "RNBQKBNRPPPP.PPP.............p..................pppp.ppprnbqkbnr" + + +@pytest.fixture +def sample_pairs(): + """Two trajectory pairs with known board strings.""" + return [(STARTING_BOARD, AFTER_E4), (STARTING_BOARD, AFTER_E4_E5)] + + +class TestTrajectoryDataset: + def test_length(self, sample_pairs): + """Dataset length matches number of trajectory pairs.""" + ds = TrajectoryDataset(sample_pairs) + assert len(ds) == 2 + + def test_empty_dataset(self): + """Empty input → empty dataset.""" + ds = TrajectoryDataset([]) + assert len(ds) == 0 + + def test_output_shapes_and_types(self, sample_pairs): + """Each item: (cur_board[64], cur_feat[14], fut_board[64], fut_feat[14]).""" + ds = TrajectoryDataset(sample_pairs) + cur_board, cur_feat, fut_board, fut_feat = ds[0] + assert cur_board.shape == (64,) and cur_board.dtype == torch.long + assert cur_feat.shape == (14,) and cur_feat.dtype == torch.float32 + assert fut_board.shape == (64,) and fut_board.dtype == torch.long + assert fut_feat.shape == (14,) and fut_feat.dtype == torch.float32 + + def test_board_values_in_piece_index_range(self, sample_pairs): + """Board values must be valid PIECE_TO_INDEX indices [0, 12].""" + ds = TrajectoryDataset(sample_pairs) + for i in range(len(ds)): + cur_board, _, fut_board, _ = ds[i] + assert (cur_board >= 0).all() and (cur_board <= 12).all(), ( + f"Pair {i}: current board has values outside [0, 12]" + ) + assert (fut_board >= 0).all() and (fut_board <= 12).all(), ( + f"Pair {i}: future board has values outside [0, 12]" + ) + + def test_current_and_future_differ(self, sample_pairs): + """Current and future boards in a pair should be different positions.""" + ds = TrajectoryDataset(sample_pairs) + cur_board, _, fut_board, _ = ds[0] + assert not torch.equal(cur_board, fut_board), ( + "Current and future boards should differ (a move was played)" + ) + + def test_known_encoding_starting_position(self): + """Starting position encoding matches PIECE_TO_INDEX manually.""" + ds = TrajectoryDataset([(STARTING_BOARD, AFTER_E4)]) + cur_board, _, _, _ = ds[0] + # First square is 'R' = 4, second 'N' = 2, etc. + expected_first_8 = [ + PIECE_TO_INDEX[c] for c in "RNBQKBNR" + ] + assert cur_board[:8].tolist() == expected_first_8 + + def test_known_encoding_future_has_moved_pawn(self): + """After e4, pawn should have moved from rank 2 to rank 4.""" + ds = TrajectoryDataset([(STARTING_BOARD, AFTER_E4)]) + _, _, fut_board, _ = ds[0] + # In AFTER_E4, index 12 (e2) is '.' = 0, index 28 (e4) is 'P' = 1 + assert fut_board[12].item() == PIECE_TO_INDEX['.'] + assert fut_board[28].item() == PIECE_TO_INDEX['P'] + + def test_features_reflect_material(self): + """Material balance feature should match board content.""" + ds = TrajectoryDataset([(STARTING_BOARD, AFTER_E4)]) + _, cur_feat, _, fut_feat = ds[0] + # Starting position: material balance = 0 (equal pieces) + expected_material = compute_features(STARTING_BOARD)[0] + assert cur_feat[0].item() == pytest.approx(expected_material) + # After e4, material is still equal (no captures) + assert fut_feat[0].item() == pytest.approx(0.0) + + def test_different_pairs_have_different_futures(self, sample_pairs): + """Different trajectory pairs should have different future boards.""" + ds = TrajectoryDataset(sample_pairs) + _, _, fut1, _ = ds[0] + _, _, fut2, _ = ds[1] + assert not torch.equal(fut1, fut2) + + def test_all_64_squares_are_valid_pieces(self): + """Every square in encoded board maps to a valid piece character.""" + valid_indices = set(PIECE_TO_INDEX.values()) + ds = TrajectoryDataset([(STARTING_BOARD, AFTER_E4)]) + cur_board, _, fut_board, _ = ds[0] + for sq in range(64): + assert cur_board[sq].item() in valid_indices + assert fut_board[sq].item() in valid_indices diff --git a/train_model.py b/train_model.py index 00334d3..e19deb5 100644 --- a/train_model.py +++ b/train_model.py @@ -1,12 +1,22 @@ -# Import required libraries -from chessformer import ChessTransformer -from chess_loader import get_dataloader +from chessformer import ChessTransformer, ChessTransformerV2 +from chess_loader import get_dataloader, get_dataloader_v2 +from grok_tracker import GrokTracker, gradfilter_ema +from model_utils import compute_loss_v2, detect_device import torch import torch.nn as nn -from torch.cuda.amp import GradScaler, autocast +import torch.nn.functional as F +from torch.amp import GradScaler, autocast import time from datetime import timedelta +# Phase 2 diffusion hyperparameters +D_DIT = 256 +DIT_NHEAD = 4 +DIT_D_HID = 512 +DIT_NLAYERS = 6 +DIFF_T = 20 +DIFF_HORIZON = 4 + # Configuration and hyperparameters single_run = True elo = 2000 @@ -20,25 +30,22 @@ INCREMENTS = num_epochs GAMMA = .9 SCALE = 1 -batch_size = int(128 / SCALE) +# Changed: 512→1024 to better utilize GPU VRAM (54% → ~80%) with 25M param model +batch_size = int(512 / SCALE) LR = {.5: 5e-4, 1: 1e-4, 2: 1e-5} LR = LR[SCALE] if SCALE in LR else 1e-5 CLIP = .1 +# orignal 512 d_model = 512 d_hid = d_model * 2 nhead = 8 nlayers = 12 +d_policy = 128 inp_ntoken = 13 out_ntoken = 2 start_time = time.time() -# Check for GPU availability -if torch.backends.mps.is_available(): - device = torch.device("mps") -elif torch.cuda.is_available(): - device = torch.device("cuda") -else: - device = torch.device("cpu") +device = detect_device() if SCALE != 1: factor = d_hid / d_model @@ -46,23 +53,77 @@ d_hid = int(d_model * factor) -def train(model: str = ""): +def _compute_loss_v1(model, boards, target, loss_fn): + output = model(boards) + return loss_fn(output, target) + + +_compute_loss_v2 = compute_loss_v2 # alias for backward compat + + +def train( + model: str = "", + patience: int = None, + model_version: str = "v1", + use_grokfast: bool = False, + grokfast_alpha: float = 0.98, + grokfast_lamb: float = 2.0, + grok_log: str | None = None, +): # Model loading or initialization + # Changed: weights_only=False + map_location for cross-device resume, supports both paths and filenames if model: - model = torch.load(f'models/{START_MODEL}') - print(f'Using Pretrained Model: {START_MODEL}') + model_path = model if '/' in model else f'models/{model}' + if model_version == 'v2': + checkpoint = torch.load(model_path, weights_only=False, map_location=device) + m = ChessTransformerV2(d_model=d_model, nhead=nhead, d_hid=d_hid, nlayers=nlayers, d_policy=d_policy, dropout=dropout) + m.load_state_dict(checkpoint['state_dict']) + model = m.to(device) + else: + model = torch.load(model_path, weights_only=False, map_location=device).to(device) + print(f'Resuming from: {model_path}') else: - model = ChessTransformer(inp_ntoken, out_ntoken, d_model, nhead, d_hid, nlayers, dropout=dropout).to(device) + if model_version == 'v2': + model = ChessTransformerV2(d_model=d_model, nhead=nhead, d_hid=d_hid, nlayers=nlayers, d_policy=d_policy, dropout=dropout).to(device) + else: + model = ChessTransformer(inp_ntoken, out_ntoken, d_model, nhead, d_hid, nlayers, dropout=dropout).to(device) + + num_params = sum(p.numel() for p in model.parameters()) + print(f'Model: {END_MODEL} ({model_version}) | Dataset: {DATASET}') + print(f'Parameters: {num_params:,} | Device: {device}') + if use_grokfast: + print(f'Grokfast: ON (alpha={grokfast_alpha}, lamb={grokfast_lamb})') + print() - print(f'Creating New Model: {END_MODEL}_best_whole.pth\nDataset: {DATASET}\n') + # Grokking tracker + Grokfast state + grok_log_path = grok_log or f'grok_{END_MODEL}.log' + tracker = GrokTracker(model, log_path=grok_log_path) + ema_grads = None # Grokfast EMA state (initialized on first backward) # Loss Function and Optimizer setup loss_fn = nn.CrossEntropyLoss() optimizer = torch.optim.AdamW(model.parameters(), lr=LR) scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(num_epochs / INCREMENTS), gamma=GAMMA) + # Changed: AMP (mixed precision) only on CUDA - not supported on MPS/CPU + use_amp = device.type == "cuda" + scaler = GradScaler("cuda") if use_amp else None # Data loaders for training and testing - dataloader, testloader = get_dataloader(DATASET, batch_size=batch_size, num_workers=0, num_pos=NUM_POS) + # Changed: num_workers=4 for parallel data loading (separate processes, no data race) + if model_version == 'v2': + dataloader, testloader = get_dataloader_v2(DATASET, batch_size=batch_size, num_workers=4, num_pos=NUM_POS) + else: + dataloader, testloader = get_dataloader(DATASET, batch_size=batch_size, num_workers=4, num_pos=NUM_POS) + + # Changed: patience tracks epochs without test loss improvement (early stopping) + no_improve = 0 + + def _batch_loss(batch_data): + if model_version == 'v2': + boards, features, from_sq, to_sq, _promo, wdl_target = [x.to(device) for x in batch_data] + return _compute_loss_v2(model, boards, features, from_sq, to_sq, wdl_target) + boards, target = batch_data[0].to(device), batch_data[1].to(device) + return _compute_loss_v1(model, boards, target, loss_fn) # Training loop best_loss = None @@ -70,51 +131,426 @@ def train(model: str = ""): for epoch in range(1, num_epochs + 1): model.train() total_loss = 0 - for batch, (boards, target) in enumerate(dataloader): - boards, target = boards.to(device), target.to(device) - output = model(boards) - loss = loss_fn(output, target) + for batch, batch_data in enumerate(dataloader): optimizer.zero_grad() - loss.backward() - if CLIP: - torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP) - optimizer.step() + if use_amp: + with autocast("cuda"): + loss = _batch_loss(batch_data) + scaler.scale(loss).backward() + scaler.unscale_(optimizer) + if use_grokfast: + ema_grads = gradfilter_ema(model, ema_grads, grokfast_alpha, grokfast_lamb) + if CLIP: + torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP) + scaler.step(optimizer) + scaler.update() + else: + loss = _batch_loss(batch_data) + loss.backward() + if use_grokfast: + ema_grads = gradfilter_ema(model, ema_grads, grokfast_alpha, grokfast_lamb) + if CLIP: + torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP) + optimizer.step() total_loss += loss.item() # Progress output - if (batch + 1) % (len(dataloader) // 10) == 0: - elapsed = str(timedelta(seconds=time.time() - start_time)) + log_interval = max(1, len(dataloader) // 10) + if (batch + 1) % log_interval == 0: + avg_loss = total_loss / (batch + 1) + elapsed = str(timedelta(seconds=int(time.time() - start_time))) print( - f'Epoch {epoch}: Batch {batch + 1} / {len(dataloader)} = {100 * (batch + 1) / len(dataloader):.{0}f}%') - print("Elapsed Time: ", elapsed) + f'Epoch {epoch}: Batch {batch + 1}/{len(dataloader)} ' + f'({100 * (batch + 1) / len(dataloader):.0f}%) ' + f'| Loss: {avg_loss:.4f} | {elapsed}') scheduler.step() + avg_train_loss = total_loss / len(dataloader) + if best_loss is None or avg_train_loss < best_loss: + best_loss = avg_train_loss # Evaluation on test dataset + model.eval() tot_test_loss = 0 with torch.no_grad(): - for batch in testloader: - boards, target = batch[0].to(device), batch[1].to(device) - output = model(boards) - loss = loss_fn(output, target) - tot_test_loss += loss.item() + for batch_data in testloader: + tot_test_loss += _batch_loss(batch_data).item() + + avg_test_loss = tot_test_loss / len(testloader) + print(f'Epoch {epoch} done | Train loss: {avg_train_loss:.4f} | Test loss: {avg_test_loss:.4f}') + tracker.log_epoch(epoch, avg_train_loss, avg_test_loss) # Save the best model - if best_test_loss is None or tot_test_loss < best_test_loss: - best_test_loss = tot_test_loss - if epoch >= min(num_epochs - 2, 3): - torch.save(model, f'models/{END_MODEL}.pth') + if best_test_loss is None or avg_test_loss < best_test_loss: + best_test_loss = avg_test_loss + no_improve = 0 + if True: + if model_version == 'v2': + torch.save({ + 'version': 'v2', + 'state_dict': model.state_dict(), + 'config': {'d_model': d_model, 'nhead': nhead, 'd_hid': d_hid, 'nlayers': nlayers, 'd_policy': d_policy}, + }, f'models/{END_MODEL}_v2.pth') + else: + torch.save(model, f'models/{END_MODEL}.pth') + print(f' -> Saved model (best test loss)') + else: + no_improve += 1 + + # Changed: early stopping - stop if test loss hasn't improved for `patience` epochs + if patience and no_improve >= patience: + print(f'\nEarly stopping: no improvement for {patience} epochs') + break + + print(f'\nBest Training Loss: {best_loss:.4f}') + print(f'Best Testing Loss: {best_test_loss:.4f}') + tracker.close() + + +def _compute_diffusion_loss(model, ns, cur_board, cur_feat, fut_board, fut_feat): + """Compute diffusion training loss (noise prediction MSE). + + 1. Encode current + future boards with frozen backbone → latents + 2. Project future latent to d_dit → x_0 + 3. Sample random timestep, add noise → x_t + 4. DiT predicts noise from (x_t, t, backbone_latent) + 5. Return MSE(predicted_noise, actual_noise) + """ + # Encode with frozen backbone + with torch.no_grad(): + backbone_latent = model.encode(cur_board, cur_feat) + future_latent = model.encode(fut_board, fut_feat) + + # Project future latent to d_dit → x_0 + x_0 = model.latent_to_dit(future_latent) + + # Diffusion forward process + B = cur_board.shape[0] + t = torch.randint(0, ns.T, (B,), device=cur_board.device) + noise = torch.randn_like(x_0) + x_t = ns.q_sample(x_0, t, noise) + + # Predict noise + predicted_noise = model.diffusion(x_t, t, backbone_latent) + + return F.mse_loss(predicted_noise, noise) + + +def train_diffusion( + backbone_model: str, + pgn_path: str, + patience: int | None = None, +): + """Phase 2: Train diffusion model on trajectory data. + + Loads a pre-trained V2 backbone, attaches a DiT diffusion model, + freezes backbone weights, and trains the diffusion components to + predict noise on future latent states. + + Args: + backbone_model: Path to pre-trained V2 model checkpoint. + pgn_path: Path to PGN file for extracting trajectories. + patience: Early stopping patience (None = disabled). + """ + from diffusion_model import ChessDiT + from noise_schedule import CosineNoiseSchedule + from trajectory_loader import get_trajectory_dataloader + + # 1. Load pre-trained V2 backbone + model_path = backbone_model if '/' in backbone_model else f'models/{backbone_model}' + checkpoint = torch.load(model_path, weights_only=False, map_location=device) + cfg = checkpoint['config'] + model = ChessTransformerV2(**cfg, dropout=0.0) + model.load_state_dict(checkpoint['state_dict']) + print(f'Loaded backbone: {model_path}') + + # 2. Create DiT + noise schedule + dit = ChessDiT( + d_dit=D_DIT, d_model=cfg['d_model'], + nhead=DIT_NHEAD, d_hid=DIT_D_HID, nlayers=DIT_NLAYERS, T=DIFF_T, + ) + ns = CosineNoiseSchedule(T=DIFF_T) + + # 3. Attach diffusion to V2 model + model.attach_diffusion(dit, ns, D_DIT) + model = model.to(device) + ns.to(device) - print(f'Best Training Loss: {best_loss / len(dataloader)}') - print(f'Best Testing Loss: {best_test_loss / len(testloader)}') + # 4. Freeze backbone, train only diffusion components + diffusion_keywords = {'diffusion', 'latent_to_dit', 'dit_to_latent'} + for name, p in model.named_parameters(): + if not any(k in name for k in diffusion_keywords): + p.requires_grad = False + + trainable = [p for p in model.parameters() if p.requires_grad] + num_trainable = sum(p.numel() for p in trainable) + num_total = sum(p.numel() for p in model.parameters()) + print(f'Parameters: {num_total:,} total, {num_trainable:,} trainable (diffusion)') + print(f'DiT config: d_dit={D_DIT}, layers={DIT_NLAYERS}, T={DIFF_T}') + print(f'Device: {device}\n') + + # 5. Optimizer + scheduler + optimizer = torch.optim.AdamW(trainable, lr=LR) + scheduler = torch.optim.lr_scheduler.StepLR( + optimizer, step_size=max(1, num_epochs // INCREMENTS), gamma=GAMMA, + ) + use_amp = device.type == "cuda" + scaler = GradScaler("cuda") if use_amp else None + + # 6. Load trajectory data + print(f'Loading trajectories from: {pgn_path} (horizon={DIFF_HORIZON})') + train_loader, test_loader = get_trajectory_dataloader( + pgn_path, horizon=DIFF_HORIZON, batch_size=batch_size, + max_trajectories=int(NUM_POS), + ) + print(f'Train batches: {len(train_loader)}, Test batches: {len(test_loader)}\n') + + # 7. Training loop + best_test_loss = None + no_improve = 0 + + for epoch in range(1, num_epochs + 1): + model.train() + total_loss = 0 + + for batch_idx, (cur_board, cur_feat, fut_board, fut_feat) in enumerate(train_loader): + cur_board = cur_board.to(device) + cur_feat = cur_feat.to(device) + fut_board = fut_board.to(device) + fut_feat = fut_feat.to(device) + + if use_amp: + with autocast("cuda"): + loss = _compute_diffusion_loss( + model, ns, cur_board, cur_feat, fut_board, fut_feat, + ) + optimizer.zero_grad() + scaler.scale(loss).backward() + if CLIP: + scaler.unscale_(optimizer) + torch.nn.utils.clip_grad_norm_(trainable, CLIP) + scaler.step(optimizer) + scaler.update() + else: + loss = _compute_diffusion_loss( + model, ns, cur_board, cur_feat, fut_board, fut_feat, + ) + optimizer.zero_grad() + loss.backward() + if CLIP: + torch.nn.utils.clip_grad_norm_(trainable, CLIP) + optimizer.step() + + total_loss += loss.item() + + log_interval = max(1, len(train_loader) // 10) + if (batch_idx + 1) % log_interval == 0: + avg = total_loss / (batch_idx + 1) + elapsed = str(timedelta(seconds=int(time.time() - start_time))) + print( + f'Epoch {epoch}: Batch {batch_idx + 1}/{len(train_loader)} ' + f'({100 * (batch_idx + 1) / len(train_loader):.0f}%) ' + f'| Diff loss: {avg:.6f} | {elapsed}' + ) + + scheduler.step() + avg_train = total_loss / len(train_loader) + + # Evaluation + model.eval() + test_loss = 0 + with torch.no_grad(): + for cur_board, cur_feat, fut_board, fut_feat in test_loader: + cur_board = cur_board.to(device) + cur_feat = cur_feat.to(device) + fut_board = fut_board.to(device) + fut_feat = fut_feat.to(device) + loss = _compute_diffusion_loss( + model, ns, cur_board, cur_feat, fut_board, fut_feat, + ) + test_loss += loss.item() + + avg_test = test_loss / len(test_loader) + print(f'Epoch {epoch} | Train diff: {avg_train:.6f} | Test diff: {avg_test:.6f}') + + # Save best model + if best_test_loss is None or test_loss < best_test_loss: + best_test_loss = test_loss + no_improve = 0 + if True: + torch.save({ + 'version': 'v2+diff', + 'state_dict': model.state_dict(), + 'config': cfg, + 'diffusion_config': { + 'd_dit': D_DIT, 'nhead': DIT_NHEAD, + 'd_hid': DIT_D_HID, 'nlayers': DIT_NLAYERS, 'T': DIFF_T, + }, + }, f'models/{END_MODEL}_v2_diff.pth') + print(f' -> Saved model (best test diff loss)') + else: + no_improve += 1 + + if patience and no_improve >= patience: + print(f'\nEarly stopping: no improvement for {patience} epochs') + break + + print(f'\nBest Test Diffusion Loss: {best_test_loss / len(test_loader):.6f}') if __name__ == '__main__': + import sys + import argparse + parser = argparse.ArgumentParser() + parser.add_argument('elo', nargs='?', type=int, default=elo, help='ELO rating filter') + # Changed: --device flag to override auto-detection (auto/cuda/mps/cpu) + parser.add_argument('--device', type=str, default='auto', choices=['auto', 'cuda', 'mps', 'cpu'], + help='Device to train on (default: auto-detect)') + # Changed: --dataset flag to override default dataset path + parser.add_argument('--dataset', type=str, default=None, + help='Path to dataset file (default: full_datasets/elo_{elo}_pos.txt)') + # Changed: --num-pos flag to override NUM_POS + parser.add_argument('--num-pos', type=float, default=NUM_POS, + help=f'Number of positions to load (default: {NUM_POS:.0f})') + # Changed: --resume loads existing model and continues training (fine-tuning) + parser.add_argument('--resume', type=str, default=None, + help='Path to existing model to continue training (e.g. models/2000_elo_pos_engine.pth)') + # Changed: --lr overrides learning rate (useful for fine-tuning with lower LR) + parser.add_argument('--lr', type=float, default=None, + help='Learning rate override (default: auto from SCALE)') + # Changed: --epochs overrides num_epochs + parser.add_argument('--epochs', type=int, default=None, + help=f'Number of epochs (default: {num_epochs})') + # Changed: --patience for early stopping - stops if test loss doesn't improve for N epochs + parser.add_argument('--patience', type=int, default=None, + help='Early stopping: stop after N epochs without test loss improvement') + # Changed: --batch-size to adjust for different VRAM sizes (1024 for 16GB, 512 for 12GB) + parser.add_argument('--batch-size', type=int, default=None, + help=f'Batch size (default: {batch_size}, lower for less VRAM)') + parser.add_argument('--model-version', type=str, default='v2', choices=['v1', 'v2'], + help='Model architecture version (default: v2)') + parser.add_argument('--grokfast', action='store_true', + help='Enable Grokfast EMA gradient filter (accelerates grokking)') + parser.add_argument('--grokfast-alpha', type=float, default=0.98, + help='Grokfast EMA decay (default: 0.98)') + parser.add_argument('--grokfast-lamb', type=float, default=2.0, + help='Grokfast amplification factor (default: 2.0)') + parser.add_argument('--grok-log', type=str, default=None, + help='Path to grokking metrics log file (default: grok_{elo}.log)') + parser.add_argument('--phase', type=int, default=1, choices=[1, 2, 3], + help='Training phase: 1=supervised, 2=diffusion, 3=selfplay (default: 1)') + parser.add_argument('--backbone-model', type=str, default=None, + help='Phase 2/3: path to pre-trained V2 model') + parser.add_argument('--pgn', type=str, default=None, + help='Phase 2: path to PGN file for trajectory extraction') + parser.add_argument('--horizon', type=int, default=DIFF_HORIZON, + help=f'Phase 2: trajectory horizon in half-moves (default: {DIFF_HORIZON})') + # Phase 3: self-play flags + parser.add_argument('--generations', type=int, default=10, + help='Phase 3: number of generate-train cycles (default: 10)') + parser.add_argument('--games-per-gen', type=int, default=100, + help='Phase 3: games to generate per generation (default: 100)') + parser.add_argument('--epochs-per-gen', type=int, default=2, + help='Phase 3: training epochs per generation (default: 2)') + parser.add_argument('--temp-schedule', type=str, default='1.5:10,1.0:25,0.3:999', + help='Phase 3: temperature schedule as "temp:ply,..." (default: 1.5:10,1.0:25,0.3:999)') + parser.add_argument('--max-moves', type=int, default=200, + help='Phase 3: max half-moves per game (default: 200)') + parser.add_argument('--buffer-size', type=int, default=5, + help='Phase 3: number of recent generations in replay buffer (default: 5)') + parser.add_argument('--mix-supervised', type=str, default=None, + help='Phase 3: path to supervised dataset for mixing') + parser.add_argument('--mix-ratio', type=float, default=0.5, + help='Phase 3: fraction of supervised data in training mix (default: 0.5)') + parser.add_argument('--eval-games', type=int, default=50, + help='Phase 3: evaluation games per generation (default: 50)') + parser.add_argument('--mcts-sims', type=int, default=0, + help='Phase 3: MCTS simulations per move (0=disabled, raw policy)') + parser.add_argument('--cpuct', type=float, default=1.25, + help='Phase 3: MCTS exploration constant (default: 1.25)') + parser.add_argument('--dirichlet-alpha', type=float, default=0.3, + help='Phase 3: Dirichlet noise alpha (default: 0.3)') + parser.add_argument('--dirichlet-epsilon', type=float, default=0.25, + help='Phase 3: Dirichlet noise weight (0=off, default: 0.25)') + parser.add_argument('--selfplay-output', type=str, default='selfplay_data', + help='Phase 3: output directory for self-play data (default: selfplay_data)') + args = parser.parse_args() + elo = args.elo + max_elo = elo + NUM_POS = args.num_pos + device = detect_device(args.device) + if args.lr is not None: + LR = args.lr + if args.epochs is not None: + num_epochs = args.epochs + INCREMENTS = num_epochs + if args.batch_size is not None: + batch_size = args.batch_size + DIFF_HORIZON = args.horizon if single_run: elo = max_elo - new_model = True - for i in range(elo, max_elo + 1, 200): - DATASET = f"full_datasets/elo_{i}_pos.txt" if FULL_SET else f"sub_datasets/elo_{i}_pos.txt" - START_MODEL = f'{i - 200}_elo_pos_engine_best_whole.pth' if not new_model else "" - END_MODEL = f'{i}_elo_pos_engine' - train(model=START_MODEL) + + if args.phase == 3: + # Phase 3: self-play training + if not args.backbone_model: + print('Error: --backbone-model is required for --phase 3') + sys.exit(1) + from selfplay_loop import selfplay_loop, SelfPlayConfig, parse_temp_schedule + config = SelfPlayConfig( + model_path=args.backbone_model, + output_dir=args.selfplay_output, + generations=args.generations, + games_per_gen=args.games_per_gen, + epochs_per_gen=args.epochs_per_gen, + batch_size=batch_size, + lr=LR, + temp_schedule=parse_temp_schedule(args.temp_schedule), + max_moves=args.max_moves, + resign_threshold=0.95, + resign_count=3, + draw_threshold=0.80, + draw_count=5, + buffer_size=args.buffer_size, + use_diffusion=False, + mix_supervised=args.mix_supervised, + mix_ratio=args.mix_ratio, + eval_games=args.eval_games, + device=args.device, + mcts_sims=args.mcts_sims, + cpuct=args.cpuct, + dirichlet_alpha=args.dirichlet_alpha, + dirichlet_epsilon=args.dirichlet_epsilon, + ) + selfplay_loop(config) + elif args.phase == 2: + # Phase 2: diffusion training + if not args.backbone_model: + print('Error: --backbone-model is required for --phase 2') + sys.exit(1) + if not args.pgn: + print('Error: --pgn is required for --phase 2') + sys.exit(1) + END_MODEL = f'{elo}_elo_pos_engine' + train_diffusion( + backbone_model=args.backbone_model, + pgn_path=args.pgn, + patience=args.patience, + ) + else: + # Phase 1: supervised training + for i in range(elo, max_elo + 1, 200): + if args.dataset: + DATASET = args.dataset + else: + DATASET = f"full_datasets/elo_{i}_pos.txt" if FULL_SET else f"sub_datasets/elo_{i}_pos.txt" + # Changed: --resume flag takes priority over START_MODEL logic + START_MODEL = args.resume if args.resume else "" + END_MODEL = f'{i}_elo_pos_engine' + train( + model=START_MODEL, patience=args.patience, + model_version=args.model_version, + use_grokfast=args.grokfast, + grokfast_alpha=args.grokfast_alpha, + grokfast_lamb=args.grokfast_lamb, + grok_log=args.grok_log, + ) diff --git a/trajectory_loader.py b/trajectory_loader.py new file mode 100644 index 0000000..fddb348 --- /dev/null +++ b/trajectory_loader.py @@ -0,0 +1,183 @@ +"""Trajectory data for Phase 2 diffusion training. + +Extracts (current_position, future_position) pairs from PGN game +trajectories. The diffusion model learns to predict future board +latents given the current position as conditioning. + +Usage: + from trajectory_loader import get_trajectory_dataloader + + train_loader, test_loader = get_trajectory_dataloader( + pgn_path="full_datasets/filtered_1500_elo.pgn", + horizon=4, + batch_size=64, + ) + +Each batch contains: + current_board: [B, 64] int — piece indices for current position + current_features: [B, 14] float — auxiliary features for current + future_board: [B, 64] int — piece indices for position H moves ahead + future_features: [B, 14] float — auxiliary features for future +""" + +import chess +import chess.pgn +import torch +from torch.utils.data import Dataset, DataLoader + +from chess_loader import PIECE_TO_INDEX, compute_features +from chess_moves_to_input_data import get_board_str + + +def extract_trajectories( + pgn_path: str, + horizon: int = 4, + min_elo: int = 1500, + max_trajectories: int | None = None, +) -> list[tuple[str, str]]: + """Extract (current, future) board string pairs from a PGN file. + + For each position P in a game, pairs it with position P+horizon + (the board state `horizon` half-moves later in the same game). + + Both board strings are 64-char representations oriented from + the current player's perspective (flipped for black via get_board_str). + + Args: + pgn_path: Path to PGN file. + horizon: Number of half-moves to look ahead (default 4 = 2 full moves). + min_elo: Minimum Elo for both players (Lichess games). + max_trajectories: Stop after collecting this many pairs. + + Returns: + List of (current_board_str, future_board_str) pairs. + """ + trajectories: list[tuple[str, str]] = [] + games_used = 0 + + with open(pgn_path) as f: + while True: + game = chess.pgn.read_game(f) + if game is None: + break + + # Filter by Elo (same criteria as pgn_to_training_data.py) + event = game.headers.get("Event", "") + is_selfplay = event.startswith("SF-SelfPlay") + + if not is_selfplay: + try: + white_elo = int(game.headers.get("WhiteElo", "0")) + black_elo = int(game.headers.get("BlackElo", "0")) + except ValueError: + continue + if white_elo < min_elo or black_elo < min_elo: + continue + if "Blitz" in event or "Bullet" in event: + continue + + # Collect all board states in the game + board = game.board() + states = [get_board_str(board, white_side=board.turn)] + for move in game.mainline_moves(): + board.push(move) + states.append(get_board_str(board, white_side=board.turn)) + + # Create trajectory pairs: (state[i], state[i + horizon]) + for i in range(len(states) - horizon): + trajectories.append((states[i], states[i + horizon])) + + games_used += 1 + if games_used % 1000 == 0: + print( + f"Games: {games_used}, trajectories: {len(trajectories)}" + ) + + if max_trajectories and len(trajectories) >= max_trajectories: + trajectories = trajectories[:max_trajectories] + break + + print( + f"Done! Games: {games_used}, trajectories: {len(trajectories)}, " + f"horizon: {horizon}" + ) + return trajectories + + +class TrajectoryDataset(Dataset): + """Board-level trajectory dataset for diffusion training. + + Each example provides current and future board states (as integer + tensors + features). During Phase 2 training, both are encoded by + the frozen V2 backbone to obtain latent representations. + """ + + def __init__(self, trajectory_pairs: list[tuple[str, str]]) -> None: + current_boards = [] + current_features = [] + future_boards = [] + future_features = [] + + for current_str, future_str in trajectory_pairs: + current_boards.append([PIECE_TO_INDEX[p] for p in current_str]) + current_features.append(compute_features(current_str)) + future_boards.append([PIECE_TO_INDEX[p] for p in future_str]) + future_features.append(compute_features(future_str)) + + self.current_boards = torch.tensor(current_boards, dtype=torch.long) + self.current_features = torch.tensor( + current_features, dtype=torch.float32 + ) + self.future_boards = torch.tensor(future_boards, dtype=torch.long) + self.future_features = torch.tensor( + future_features, dtype=torch.float32 + ) + + def __len__(self) -> int: + return len(self.current_boards) + + def __getitem__( + self, idx: int + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return ( + self.current_boards[idx], + self.current_features[idx], + self.future_boards[idx], + self.future_features[idx], + ) + + +def get_trajectory_dataloader( + pgn_path: str, + horizon: int = 4, + batch_size: int = 64, + min_elo: int = 1500, + max_trajectories: int | None = None, + num_workers: int = 0, +) -> tuple[DataLoader, DataLoader]: + """Create train/test dataloaders from PGN trajectory pairs.""" + pairs = extract_trajectories( + pgn_path, horizon, min_elo, max_trajectories + ) + dataset = TrajectoryDataset(pairs) + + test_len = min(5000, int(len(dataset) * 0.1)) + train_set, test_set = torch.utils.data.random_split( + dataset, [len(dataset) - test_len, test_len] + ) + + train_loader = DataLoader( + train_set, + batch_size=batch_size, + shuffle=True, + pin_memory=True, + num_workers=num_workers, + ) + test_loader = DataLoader( + test_set, + batch_size=batch_size, + shuffle=True, + pin_memory=True, + num_workers=num_workers, + ) + return train_loader, test_loader diff --git a/transformer.py b/transformer.py new file mode 100644 index 0000000..ffd81e2 --- /dev/null +++ b/transformer.py @@ -0,0 +1,10 @@ +"""Compatibility shim for older pickled models. + +The released model was saved with a module path named "transformer". +This file re-exports the current implementation from chessformer.py +so torch.load can resolve the class during unpickling. +""" + +from chessformer import ChessTransformer, PositionalEncoding # re-export for torch.load + +__all__ = ["ChessTransformer", "PositionalEncoding"] diff --git a/uv.lock b/uv.lock new file mode 100644 index 0000000..5392968 --- /dev/null +++ b/uv.lock @@ -0,0 +1,574 @@ +version = 1 +revision = 3 +requires-python = ">=3.12" +resolution-markers = [ + "(python_full_version >= '3.15' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')", + "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'", +] + +[[package]] +name = "chess" +version = "1.11.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/93/09/7d04d7581ae3bb8b598017941781bceb7959dd1b13e3ebf7b6a2cd843bc9/chess-1.11.2.tar.gz", hash = "sha256:a8b43e5678fdb3000695bdaa573117ad683761e5ca38e591c4826eba6d25bb39", size = 6131385, upload-time = "2025-02-25T19:10:27.328Z" } + +[[package]] +name = "chessformer" +version = "0.1.0" +source = { editable = "." } +dependencies = [ + { name = "numpy" }, + { name = "pygame" }, + { name = "python-chess" }, + { name = "torch", version = "2.9.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'" }, + { name = "torch", version = "2.9.1+rocm7.2.0.lw.git7e1940d4", source = { registry = "wheelies" }, marker = "(python_full_version >= '3.15' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')" }, +] + +[package.dev-dependencies] +dev = [ + { name = "pytest" }, +] + +[package.metadata] +requires-dist = [ + { name = "numpy", specifier = ">=1.21.0" }, + { name = "pygame", specifier = ">=2.6.1" }, + { name = "python-chess", specifier = ">=1.999" }, + { name = "torch", specifier = ">=2.0.0" }, +] + +[package.metadata.requires-dev] +dev = [{ name = "pytest", specifier = ">=9.0.2" }] + +[[package]] +name = "colorama" +version = "0.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, +] + +[[package]] +name = "filelock" +version = "3.24.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/02/a8/dae62680be63cbb3ff87cfa2f51cf766269514ea5488479d42fec5aa6f3a/filelock-3.24.2.tar.gz", hash = "sha256:c22803117490f156e59fafce621f0550a7a853e2bbf4f87f112b11d469b6c81b", size = 37601, upload-time = "2026-02-16T02:50:45.614Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/04/a94ebfb4eaaa08db56725a40de2887e95de4e8641b9e902c311bfa00aa39/filelock-3.24.2-py3-none-any.whl", hash = "sha256:667d7dc0b7d1e1064dd5f8f8e80bdac157a6482e8d2e02cd16fd3b6b33bd6556", size = 24152, upload-time = "2026-02-16T02:50:44Z" }, +] + +[[package]] +name = "fsspec" +version = "2026.2.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/51/7c/f60c259dcbf4f0c47cc4ddb8f7720d2dcdc8888c8e5ad84c73ea4531cc5b/fsspec-2026.2.0.tar.gz", hash = "sha256:6544e34b16869f5aacd5b90bdf1a71acb37792ea3ddf6125ee69a22a53fb8bff", size = 313441, upload-time = "2026-02-05T21:50:53.743Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e6/ab/fb21f4c939bb440104cc2b396d3be1d9b7a9fd3c6c2a53d98c45b3d7c954/fsspec-2026.2.0-py3-none-any.whl", hash = "sha256:98de475b5cb3bd66bedd5c4679e87b4fdfe1a3bf4d707b151b3c07e58c9a2437", size = 202505, upload-time = "2026-02-05T21:50:51.819Z" }, +] + +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, +] + +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/72/147da192e38635ada20e0a2e1a51cf8823d2119ce8883f7053879c2199b5/markupsafe-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d53197da72cc091b024dd97249dfc7794d6a56530370992a5e1a08983ad9230e", size = 11615, upload-time = "2025-09-27T18:36:30.854Z" }, + { url = "https://files.pythonhosted.org/packages/9a/81/7e4e08678a1f98521201c3079f77db69fb552acd56067661f8c2f534a718/markupsafe-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1872df69a4de6aead3491198eaf13810b565bdbeec3ae2dc8780f14458ec73ce", size = 12020, upload-time = "2025-09-27T18:36:31.971Z" }, + { url = "https://files.pythonhosted.org/packages/1e/2c/799f4742efc39633a1b54a92eec4082e4f815314869865d876824c257c1e/markupsafe-3.0.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3a7e8ae81ae39e62a41ec302f972ba6ae23a5c5396c8e60113e9066ef893da0d", size = 24332, upload-time = "2025-09-27T18:36:32.813Z" }, + { url = "https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" }, + { url = "https://files.pythonhosted.org/packages/2c/54/887f3092a85238093a0b2154bd629c89444f395618842e8b0c41783898ea/markupsafe-3.0.3-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:94c6f0bb423f739146aec64595853541634bde58b2135f27f61c1ffd1cd4d16a", size = 21962, upload-time = "2025-09-27T18:36:35.099Z" }, + { url = "https://files.pythonhosted.org/packages/c9/2f/336b8c7b6f4a4d95e91119dc8521402461b74a485558d8f238a68312f11c/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:be8813b57049a7dc738189df53d69395eba14fb99345e0a5994914a3864c8a4b", size = 23760, upload-time = "2025-09-27T18:36:36.001Z" }, + { url = "https://files.pythonhosted.org/packages/32/43/67935f2b7e4982ffb50a4d169b724d74b62a3964bc1a9a527f5ac4f1ee2b/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:83891d0e9fb81a825d9a6d61e3f07550ca70a076484292a70fde82c4b807286f", size = 21529, upload-time = "2025-09-27T18:36:36.906Z" }, + { url = "https://files.pythonhosted.org/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" }, + { url = "https://files.pythonhosted.org/packages/2f/e1/78ee7a023dac597a5825441ebd17170785a9dab23de95d2c7508ade94e0e/markupsafe-3.0.3-cp312-cp312-win32.whl", hash = "sha256:d88b440e37a16e651bda4c7c2b930eb586fd15ca7406cb39e211fcff3bf3017d", size = 14540, upload-time = "2025-09-27T18:36:38.761Z" }, + { url = "https://files.pythonhosted.org/packages/aa/5b/bec5aa9bbbb2c946ca2733ef9c4ca91c91b6a24580193e891b5f7dbe8e1e/markupsafe-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:26a5784ded40c9e318cfc2bdb30fe164bdb8665ded9cd64d500a34fb42067b1c", size = 15105, upload-time = "2025-09-27T18:36:39.701Z" }, + { url = "https://files.pythonhosted.org/packages/e5/f1/216fc1bbfd74011693a4fd837e7026152e89c4bcf3e77b6692fba9923123/markupsafe-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:35add3b638a5d900e807944a078b51922212fb3dedb01633a8defc4b01a3c85f", size = 13906, upload-time = "2025-09-27T18:36:40.689Z" }, + { url = "https://files.pythonhosted.org/packages/38/2f/907b9c7bbba283e68f20259574b13d005c121a0fa4c175f9bed27c4597ff/markupsafe-3.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:e1cf1972137e83c5d4c136c43ced9ac51d0e124706ee1c8aa8532c1287fa8795", size = 11622, upload-time = "2025-09-27T18:36:41.777Z" }, + { url = "https://files.pythonhosted.org/packages/9c/d9/5f7756922cdd676869eca1c4e3c0cd0df60ed30199ffd775e319089cb3ed/markupsafe-3.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:116bb52f642a37c115f517494ea5feb03889e04df47eeff5b130b1808ce7c219", size = 12029, upload-time = "2025-09-27T18:36:43.257Z" }, + { url = "https://files.pythonhosted.org/packages/00/07/575a68c754943058c78f30db02ee03a64b3c638586fba6a6dd56830b30a3/markupsafe-3.0.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:133a43e73a802c5562be9bbcd03d090aa5a1fe899db609c29e8c8d815c5f6de6", size = 24374, upload-time = "2025-09-27T18:36:44.508Z" }, + { url = "https://files.pythonhosted.org/packages/a9/21/9b05698b46f218fc0e118e1f8168395c65c8a2c750ae2bab54fc4bd4e0e8/markupsafe-3.0.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ccfcd093f13f0f0b7fdd0f198b90053bf7b2f02a3927a30e63f3ccc9df56b676", size = 22980, upload-time = "2025-09-27T18:36:45.385Z" }, + { url = "https://files.pythonhosted.org/packages/7f/71/544260864f893f18b6827315b988c146b559391e6e7e8f7252839b1b846a/markupsafe-3.0.3-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:509fa21c6deb7a7a273d629cf5ec029bc209d1a51178615ddf718f5918992ab9", size = 21990, upload-time = "2025-09-27T18:36:46.916Z" }, + { url = "https://files.pythonhosted.org/packages/c2/28/b50fc2f74d1ad761af2f5dcce7492648b983d00a65b8c0e0cb457c82ebbe/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:a4afe79fb3de0b7097d81da19090f4df4f8d3a2b3adaa8764138aac2e44f3af1", size = 23784, upload-time = "2025-09-27T18:36:47.884Z" }, + { url = "https://files.pythonhosted.org/packages/ed/76/104b2aa106a208da8b17a2fb72e033a5a9d7073c68f7e508b94916ed47a9/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:795e7751525cae078558e679d646ae45574b47ed6e7771863fcc079a6171a0fc", size = 21588, upload-time = "2025-09-27T18:36:48.82Z" }, + { url = "https://files.pythonhosted.org/packages/b5/99/16a5eb2d140087ebd97180d95249b00a03aa87e29cc224056274f2e45fd6/markupsafe-3.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:8485f406a96febb5140bfeca44a73e3ce5116b2501ac54fe953e488fb1d03b12", size = 23041, upload-time = "2025-09-27T18:36:49.797Z" }, + { url = "https://files.pythonhosted.org/packages/19/bc/e7140ed90c5d61d77cea142eed9f9c303f4c4806f60a1044c13e3f1471d0/markupsafe-3.0.3-cp313-cp313-win32.whl", hash = "sha256:bdd37121970bfd8be76c5fb069c7751683bdf373db1ed6c010162b2a130248ed", size = 14543, upload-time = "2025-09-27T18:36:51.584Z" }, + { url = "https://files.pythonhosted.org/packages/05/73/c4abe620b841b6b791f2edc248f556900667a5a1cf023a6646967ae98335/markupsafe-3.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:9a1abfdc021a164803f4d485104931fb8f8c1efd55bc6b748d2f5774e78b62c5", size = 15113, upload-time = "2025-09-27T18:36:52.537Z" }, + { url = "https://files.pythonhosted.org/packages/f0/3a/fa34a0f7cfef23cf9500d68cb7c32dd64ffd58a12b09225fb03dd37d5b80/markupsafe-3.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:7e68f88e5b8799aa49c85cd116c932a1ac15caaa3f5db09087854d218359e485", size = 13911, upload-time = "2025-09-27T18:36:53.513Z" }, + { url = "https://files.pythonhosted.org/packages/e4/d7/e05cd7efe43a88a17a37b3ae96e79a19e846f3f456fe79c57ca61356ef01/markupsafe-3.0.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:218551f6df4868a8d527e3062d0fb968682fe92054e89978594c28e642c43a73", size = 11658, upload-time = "2025-09-27T18:36:54.819Z" }, + { url = "https://files.pythonhosted.org/packages/99/9e/e412117548182ce2148bdeacdda3bb494260c0b0184360fe0d56389b523b/markupsafe-3.0.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:3524b778fe5cfb3452a09d31e7b5adefeea8c5be1d43c4f810ba09f2ceb29d37", size = 12066, upload-time = "2025-09-27T18:36:55.714Z" }, + { url = "https://files.pythonhosted.org/packages/bc/e6/fa0ffcda717ef64a5108eaa7b4f5ed28d56122c9a6d70ab8b72f9f715c80/markupsafe-3.0.3-cp313-cp313t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4e885a3d1efa2eadc93c894a21770e4bc67899e3543680313b09f139e149ab19", size = 25639, upload-time = "2025-09-27T18:36:56.908Z" }, + { url = "https://files.pythonhosted.org/packages/96/ec/2102e881fe9d25fc16cb4b25d5f5cde50970967ffa5dddafdb771237062d/markupsafe-3.0.3-cp313-cp313t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8709b08f4a89aa7586de0aadc8da56180242ee0ada3999749b183aa23df95025", size = 23569, upload-time = "2025-09-27T18:36:57.913Z" }, + { url = "https://files.pythonhosted.org/packages/4b/30/6f2fce1f1f205fc9323255b216ca8a235b15860c34b6798f810f05828e32/markupsafe-3.0.3-cp313-cp313t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:b8512a91625c9b3da6f127803b166b629725e68af71f8184ae7e7d54686a56d6", size = 23284, upload-time = "2025-09-27T18:36:58.833Z" }, + { url = "https://files.pythonhosted.org/packages/58/47/4a0ccea4ab9f5dcb6f79c0236d954acb382202721e704223a8aafa38b5c8/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:9b79b7a16f7fedff2495d684f2b59b0457c3b493778c9eed31111be64d58279f", size = 24801, upload-time = "2025-09-27T18:36:59.739Z" }, + { url = "https://files.pythonhosted.org/packages/6a/70/3780e9b72180b6fecb83a4814d84c3bf4b4ae4bf0b19c27196104149734c/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_riscv64.whl", hash = "sha256:12c63dfb4a98206f045aa9563db46507995f7ef6d83b2f68eda65c307c6829eb", size = 22769, upload-time = "2025-09-27T18:37:00.719Z" }, + { url = "https://files.pythonhosted.org/packages/98/c5/c03c7f4125180fc215220c035beac6b9cb684bc7a067c84fc69414d315f5/markupsafe-3.0.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:8f71bc33915be5186016f675cd83a1e08523649b0e33efdb898db577ef5bb009", size = 23642, upload-time = "2025-09-27T18:37:01.673Z" }, + { url = "https://files.pythonhosted.org/packages/80/d6/2d1b89f6ca4bff1036499b1e29a1d02d282259f3681540e16563f27ebc23/markupsafe-3.0.3-cp313-cp313t-win32.whl", hash = "sha256:69c0b73548bc525c8cb9a251cddf1931d1db4d2258e9599c28c07ef3580ef354", size = 14612, upload-time = "2025-09-27T18:37:02.639Z" }, + { url = "https://files.pythonhosted.org/packages/2b/98/e48a4bfba0a0ffcf9925fe2d69240bfaa19c6f7507b8cd09c70684a53c1e/markupsafe-3.0.3-cp313-cp313t-win_amd64.whl", hash = "sha256:1b4b79e8ebf6b55351f0d91fe80f893b4743f104bff22e90697db1590e47a218", size = 15200, upload-time = "2025-09-27T18:37:03.582Z" }, + { url = "https://files.pythonhosted.org/packages/0e/72/e3cc540f351f316e9ed0f092757459afbc595824ca724cbc5a5d4263713f/markupsafe-3.0.3-cp313-cp313t-win_arm64.whl", hash = "sha256:ad2cf8aa28b8c020ab2fc8287b0f823d0a7d8630784c31e9ee5edea20f406287", size = 13973, upload-time = "2025-09-27T18:37:04.929Z" }, + { url = "https://files.pythonhosted.org/packages/33/8a/8e42d4838cd89b7dde187011e97fe6c3af66d8c044997d2183fbd6d31352/markupsafe-3.0.3-cp314-cp314-macosx_10_13_x86_64.whl", hash = "sha256:eaa9599de571d72e2daf60164784109f19978b327a3910d3e9de8c97b5b70cfe", size = 11619, upload-time = "2025-09-27T18:37:06.342Z" }, + { url = "https://files.pythonhosted.org/packages/b5/64/7660f8a4a8e53c924d0fa05dc3a55c9cee10bbd82b11c5afb27d44b096ce/markupsafe-3.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:c47a551199eb8eb2121d4f0f15ae0f923d31350ab9280078d1e5f12b249e0026", size = 12029, upload-time = "2025-09-27T18:37:07.213Z" }, + { url = "https://files.pythonhosted.org/packages/da/ef/e648bfd021127bef5fa12e1720ffed0c6cbb8310c8d9bea7266337ff06de/markupsafe-3.0.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f34c41761022dd093b4b6896d4810782ffbabe30f2d443ff5f083e0cbbb8c737", size = 24408, upload-time = "2025-09-27T18:37:09.572Z" }, + { url = "https://files.pythonhosted.org/packages/41/3c/a36c2450754618e62008bf7435ccb0f88053e07592e6028a34776213d877/markupsafe-3.0.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:457a69a9577064c05a97c41f4e65148652db078a3a509039e64d3467b9e7ef97", size = 23005, upload-time = "2025-09-27T18:37:10.58Z" }, + { url = "https://files.pythonhosted.org/packages/bc/20/b7fdf89a8456b099837cd1dc21974632a02a999ec9bf7ca3e490aacd98e7/markupsafe-3.0.3-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e8afc3f2ccfa24215f8cb28dcf43f0113ac3c37c2f0f0806d8c70e4228c5cf4d", size = 22048, upload-time = "2025-09-27T18:37:11.547Z" }, + { url = "https://files.pythonhosted.org/packages/9a/a7/591f592afdc734f47db08a75793a55d7fbcc6902a723ae4cfbab61010cc5/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ec15a59cf5af7be74194f7ab02d0f59a62bdcf1a537677ce67a2537c9b87fcda", size = 23821, upload-time = "2025-09-27T18:37:12.48Z" }, + { url = "https://files.pythonhosted.org/packages/7d/33/45b24e4f44195b26521bc6f1a82197118f74df348556594bd2262bda1038/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:0eb9ff8191e8498cca014656ae6b8d61f39da5f95b488805da4bb029cccbfbaf", size = 21606, upload-time = "2025-09-27T18:37:13.485Z" }, + { url = "https://files.pythonhosted.org/packages/ff/0e/53dfaca23a69fbfbbf17a4b64072090e70717344c52eaaaa9c5ddff1e5f0/markupsafe-3.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:2713baf880df847f2bece4230d4d094280f4e67b1e813eec43b4c0e144a34ffe", size = 23043, upload-time = "2025-09-27T18:37:14.408Z" }, + { url = "https://files.pythonhosted.org/packages/46/11/f333a06fc16236d5238bfe74daccbca41459dcd8d1fa952e8fbd5dccfb70/markupsafe-3.0.3-cp314-cp314-win32.whl", hash = "sha256:729586769a26dbceff69f7a7dbbf59ab6572b99d94576a5592625d5b411576b9", size = 14747, upload-time = "2025-09-27T18:37:15.36Z" }, + { url = "https://files.pythonhosted.org/packages/28/52/182836104b33b444e400b14f797212f720cbc9ed6ba34c800639d154e821/markupsafe-3.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:bdc919ead48f234740ad807933cdf545180bfbe9342c2bb451556db2ed958581", size = 15341, upload-time = "2025-09-27T18:37:16.496Z" }, + { url = "https://files.pythonhosted.org/packages/6f/18/acf23e91bd94fd7b3031558b1f013adfa21a8e407a3fdb32745538730382/markupsafe-3.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:5a7d5dc5140555cf21a6fefbdbf8723f06fcd2f63ef108f2854de715e4422cb4", size = 14073, upload-time = "2025-09-27T18:37:17.476Z" }, + { url = "https://files.pythonhosted.org/packages/3c/f0/57689aa4076e1b43b15fdfa646b04653969d50cf30c32a102762be2485da/markupsafe-3.0.3-cp314-cp314t-macosx_10_13_x86_64.whl", hash = "sha256:1353ef0c1b138e1907ae78e2f6c63ff67501122006b0f9abad68fda5f4ffc6ab", size = 11661, upload-time = "2025-09-27T18:37:18.453Z" }, + { url = "https://files.pythonhosted.org/packages/89/c3/2e67a7ca217c6912985ec766c6393b636fb0c2344443ff9d91404dc4c79f/markupsafe-3.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:1085e7fbddd3be5f89cc898938f42c0b3c711fdcb37d75221de2666af647c175", size = 12069, upload-time = "2025-09-27T18:37:19.332Z" }, + { url = "https://files.pythonhosted.org/packages/f0/00/be561dce4e6ca66b15276e184ce4b8aec61fe83662cce2f7d72bd3249d28/markupsafe-3.0.3-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1b52b4fb9df4eb9ae465f8d0c228a00624de2334f216f178a995ccdcf82c4634", size = 25670, upload-time = "2025-09-27T18:37:20.245Z" }, + { url = "https://files.pythonhosted.org/packages/50/09/c419f6f5a92e5fadde27efd190eca90f05e1261b10dbd8cbcb39cd8ea1dc/markupsafe-3.0.3-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fed51ac40f757d41b7c48425901843666a6677e3e8eb0abcff09e4ba6e664f50", size = 23598, upload-time = "2025-09-27T18:37:21.177Z" }, + { url = "https://files.pythonhosted.org/packages/22/44/a0681611106e0b2921b3033fc19bc53323e0b50bc70cffdd19f7d679bb66/markupsafe-3.0.3-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:f190daf01f13c72eac4efd5c430a8de82489d9cff23c364c3ea822545032993e", size = 23261, upload-time = "2025-09-27T18:37:22.167Z" }, + { url = "https://files.pythonhosted.org/packages/5f/57/1b0b3f100259dc9fffe780cfb60d4be71375510e435efec3d116b6436d43/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:e56b7d45a839a697b5eb268c82a71bd8c7f6c94d6fd50c3d577fa39a9f1409f5", size = 24835, upload-time = "2025-09-27T18:37:23.296Z" }, + { url = "https://files.pythonhosted.org/packages/26/6a/4bf6d0c97c4920f1597cc14dd720705eca0bf7c787aebc6bb4d1bead5388/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:f3e98bb3798ead92273dc0e5fd0f31ade220f59a266ffd8a4f6065e0a3ce0523", size = 22733, upload-time = "2025-09-27T18:37:24.237Z" }, + { url = "https://files.pythonhosted.org/packages/14/c7/ca723101509b518797fedc2fdf79ba57f886b4aca8a7d31857ba3ee8281f/markupsafe-3.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5678211cb9333a6468fb8d8be0305520aa073f50d17f089b5b4b477ea6e67fdc", size = 23672, upload-time = "2025-09-27T18:37:25.271Z" }, + { url = "https://files.pythonhosted.org/packages/fb/df/5bd7a48c256faecd1d36edc13133e51397e41b73bb77e1a69deab746ebac/markupsafe-3.0.3-cp314-cp314t-win32.whl", hash = "sha256:915c04ba3851909ce68ccc2b8e2cd691618c4dc4c4232fb7982bca3f41fd8c3d", size = 14819, upload-time = "2025-09-27T18:37:26.285Z" }, + { url = "https://files.pythonhosted.org/packages/1a/8a/0402ba61a2f16038b48b39bccca271134be00c5c9f0f623208399333c448/markupsafe-3.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4faffd047e07c38848ce017e8725090413cd80cbc23d86e55c587bf979e579c9", size = 15426, upload-time = "2025-09-27T18:37:27.316Z" }, + { url = "https://files.pythonhosted.org/packages/70/bc/6f1c2f612465f5fa89b95bead1f44dcb607670fd42891d8fdcd5d039f4f4/markupsafe-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:32001d6a8fc98c8cb5c947787c5d08b0a50663d139f1305bac5885d98d9b40fa", size = 14146, upload-time = "2025-09-27T18:37:28.327Z" }, +] + +[[package]] +name = "mpmath" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, +] + +[[package]] +name = "networkx" +version = "3.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, +] + +[[package]] +name = "numpy" +version = "2.4.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/57/fd/0005efbd0af48e55eb3c7208af93f2862d4b1a56cd78e84309a2d959208d/numpy-2.4.2.tar.gz", hash = "sha256:659a6107e31a83c4e33f763942275fd278b21d095094044eb35569e86a21ddae", size = 20723651, upload-time = "2026-01-31T23:13:10.135Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/51/6e/6f394c9c77668153e14d4da83bcc247beb5952f6ead7699a1a2992613bea/numpy-2.4.2-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:21982668592194c609de53ba4933a7471880ccbaadcc52352694a59ecc860b3a", size = 16667963, upload-time = "2026-01-31T23:10:52.147Z" }, + { url = "https://files.pythonhosted.org/packages/1f/f8/55483431f2b2fd015ae6ed4fe62288823ce908437ed49db5a03d15151678/numpy-2.4.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40397bda92382fcec844066efb11f13e1c9a3e2a8e8f318fb72ed8b6db9f60f1", size = 14693571, upload-time = "2026-01-31T23:10:54.789Z" }, + { url = "https://files.pythonhosted.org/packages/2f/20/18026832b1845cdc82248208dd929ca14c9d8f2bac391f67440707fff27c/numpy-2.4.2-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:b3a24467af63c67829bfaa61eecf18d5432d4f11992688537be59ecd6ad32f5e", size = 5203469, upload-time = "2026-01-31T23:10:57.343Z" }, + { url = "https://files.pythonhosted.org/packages/7d/33/2eb97c8a77daaba34eaa3fa7241a14ac5f51c46a6bd5911361b644c4a1e2/numpy-2.4.2-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:805cc8de9fd6e7a22da5aed858e0ab16be5a4db6c873dde1d7451c541553aa27", size = 6550820, upload-time = "2026-01-31T23:10:59.429Z" }, + { url = "https://files.pythonhosted.org/packages/b1/91/b97fdfd12dc75b02c44e26c6638241cc004d4079a0321a69c62f51470c4c/numpy-2.4.2-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6d82351358ffbcdcd7b686b90742a9b86632d6c1c051016484fa0b326a0a1548", size = 15663067, upload-time = "2026-01-31T23:11:01.291Z" }, + { url = "https://files.pythonhosted.org/packages/f5/c6/a18e59f3f0b8071cc85cbc8d80cd02d68aa9710170b2553a117203d46936/numpy-2.4.2-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9e35d3e0144137d9fdae62912e869136164534d64a169f86438bc9561b6ad49f", size = 16619782, upload-time = "2026-01-31T23:11:03.669Z" }, + { url = "https://files.pythonhosted.org/packages/b7/83/9751502164601a79e18847309f5ceec0b1446d7b6aa12305759b72cf98b2/numpy-2.4.2-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:adb6ed2ad29b9e15321d167d152ee909ec73395901b70936f029c3bc6d7f4460", size = 17013128, upload-time = "2026-01-31T23:11:05.913Z" }, + { url = "https://files.pythonhosted.org/packages/61/c4/c4066322256ec740acc1c8923a10047818691d2f8aec254798f3dd90f5f2/numpy-2.4.2-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:8906e71fd8afcb76580404e2a950caef2685df3d2a57fe82a86ac8d33cc007ba", size = 18345324, upload-time = "2026-01-31T23:11:08.248Z" }, + { url = "https://files.pythonhosted.org/packages/ab/af/6157aa6da728fa4525a755bfad486ae7e3f76d4c1864138003eb84328497/numpy-2.4.2-cp312-cp312-win32.whl", hash = "sha256:ec055f6dae239a6299cace477b479cca2fc125c5675482daf1dd886933a1076f", size = 5960282, upload-time = "2026-01-31T23:11:10.497Z" }, + { url = "https://files.pythonhosted.org/packages/92/0f/7ceaaeaacb40567071e94dbf2c9480c0ae453d5bb4f52bea3892c39dc83c/numpy-2.4.2-cp312-cp312-win_amd64.whl", hash = "sha256:209fae046e62d0ce6435fcfe3b1a10537e858249b3d9b05829e2a05218296a85", size = 12314210, upload-time = "2026-01-31T23:11:12.176Z" }, + { url = "https://files.pythonhosted.org/packages/2f/a3/56c5c604fae6dd40fa2ed3040d005fca97e91bd320d232ac9931d77ba13c/numpy-2.4.2-cp312-cp312-win_arm64.whl", hash = "sha256:fbde1b0c6e81d56f5dccd95dd4a711d9b95df1ae4009a60887e56b27e8d903fa", size = 10220171, upload-time = "2026-01-31T23:11:14.684Z" }, + { url = "https://files.pythonhosted.org/packages/a1/22/815b9fe25d1d7ae7d492152adbc7226d3eff731dffc38fe970589fcaaa38/numpy-2.4.2-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:25f2059807faea4b077a2b6837391b5d830864b3543627f381821c646f31a63c", size = 16663696, upload-time = "2026-01-31T23:11:17.516Z" }, + { url = "https://files.pythonhosted.org/packages/09/f0/817d03a03f93ba9c6c8993de509277d84e69f9453601915e4a69554102a1/numpy-2.4.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:bd3a7a9f5847d2fb8c2c6d1c862fa109c31a9abeca1a3c2bd5a64572955b2979", size = 14688322, upload-time = "2026-01-31T23:11:19.883Z" }, + { url = "https://files.pythonhosted.org/packages/da/b4/f805ab79293c728b9a99438775ce51885fd4f31b76178767cfc718701a39/numpy-2.4.2-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:8e4549f8a3c6d13d55041925e912bfd834285ef1dd64d6bc7d542583355e2e98", size = 5198157, upload-time = "2026-01-31T23:11:22.375Z" }, + { url = "https://files.pythonhosted.org/packages/74/09/826e4289844eccdcd64aac27d13b0fd3f32039915dd5b9ba01baae1f436c/numpy-2.4.2-cp313-cp313-macosx_14_0_x86_64.whl", hash = "sha256:aea4f66ff44dfddf8c2cffd66ba6538c5ec67d389285292fe428cb2c738c8aef", size = 6546330, upload-time = "2026-01-31T23:11:23.958Z" }, + { url = "https://files.pythonhosted.org/packages/19/fb/cbfdbfa3057a10aea5422c558ac57538e6acc87ec1669e666d32ac198da7/numpy-2.4.2-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c3cd545784805de05aafe1dde61752ea49a359ccba9760c1e5d1c88a93bbf2b7", size = 15660968, upload-time = "2026-01-31T23:11:25.713Z" }, + { url = "https://files.pythonhosted.org/packages/04/dc/46066ce18d01645541f0186877377b9371b8fa8017fa8262002b4ef22612/numpy-2.4.2-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d0d9b7c93578baafcbc5f0b83eaf17b79d345c6f36917ba0c67f45226911d499", size = 16607311, upload-time = "2026-01-31T23:11:28.117Z" }, + { url = "https://files.pythonhosted.org/packages/14/d9/4b5adfc39a43fa6bf918c6d544bc60c05236cc2f6339847fc5b35e6cb5b0/numpy-2.4.2-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:f74f0f7779cc7ae07d1810aab8ac6b1464c3eafb9e283a40da7309d5e6e48fbb", size = 17012850, upload-time = "2026-01-31T23:11:30.888Z" }, + { url = "https://files.pythonhosted.org/packages/b7/20/adb6e6adde6d0130046e6fdfb7675cc62bc2f6b7b02239a09eb58435753d/numpy-2.4.2-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:c7ac672d699bf36275c035e16b65539931347d68b70667d28984c9fb34e07fa7", size = 18334210, upload-time = "2026-01-31T23:11:33.214Z" }, + { url = "https://files.pythonhosted.org/packages/78/0e/0a73b3dff26803a8c02baa76398015ea2a5434d9b8265a7898a6028c1591/numpy-2.4.2-cp313-cp313-win32.whl", hash = "sha256:8e9afaeb0beff068b4d9cd20d322ba0ee1cecfb0b08db145e4ab4dd44a6b5110", size = 5958199, upload-time = "2026-01-31T23:11:35.385Z" }, + { url = "https://files.pythonhosted.org/packages/43/bc/6352f343522fcb2c04dbaf94cb30cca6fd32c1a750c06ad6231b4293708c/numpy-2.4.2-cp313-cp313-win_amd64.whl", hash = "sha256:7df2de1e4fba69a51c06c28f5a3de36731eb9639feb8e1cf7e4a7b0daf4cf622", size = 12310848, upload-time = "2026-01-31T23:11:38.001Z" }, + { url = "https://files.pythonhosted.org/packages/6e/8d/6da186483e308da5da1cc6918ce913dcfe14ffde98e710bfeff2a6158d4e/numpy-2.4.2-cp313-cp313-win_arm64.whl", hash = "sha256:0fece1d1f0a89c16b03442eae5c56dc0be0c7883b5d388e0c03f53019a4bfd71", size = 10221082, upload-time = "2026-01-31T23:11:40.392Z" }, + { url = "https://files.pythonhosted.org/packages/25/a1/9510aa43555b44781968935c7548a8926274f815de42ad3997e9e83680dd/numpy-2.4.2-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:5633c0da313330fd20c484c78cdd3f9b175b55e1a766c4a174230c6b70ad8262", size = 14815866, upload-time = "2026-01-31T23:11:42.495Z" }, + { url = "https://files.pythonhosted.org/packages/36/30/6bbb5e76631a5ae46e7923dd16ca9d3f1c93cfa8d4ed79a129814a9d8db3/numpy-2.4.2-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:d9f64d786b3b1dd742c946c42d15b07497ed14af1a1f3ce840cce27daa0ce913", size = 5325631, upload-time = "2026-01-31T23:11:44.7Z" }, + { url = "https://files.pythonhosted.org/packages/46/00/3a490938800c1923b567b3a15cd17896e68052e2145d8662aaf3e1ffc58f/numpy-2.4.2-cp313-cp313t-macosx_14_0_x86_64.whl", hash = "sha256:b21041e8cb6a1eb5312dd1d2f80a94d91efffb7a06b70597d44f1bd2dfc315ab", size = 6646254, upload-time = "2026-01-31T23:11:46.341Z" }, + { url = "https://files.pythonhosted.org/packages/d3/e9/fac0890149898a9b609caa5af7455a948b544746e4b8fe7c212c8edd71f8/numpy-2.4.2-cp313-cp313t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:00ab83c56211a1d7c07c25e3217ea6695e50a3e2f255053686b081dc0b091a82", size = 15720138, upload-time = "2026-01-31T23:11:48.082Z" }, + { url = "https://files.pythonhosted.org/packages/ea/5c/08887c54e68e1e28df53709f1893ce92932cc6f01f7c3d4dc952f61ffd4e/numpy-2.4.2-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2fb882da679409066b4603579619341c6d6898fc83a8995199d5249f986e8e8f", size = 16655398, upload-time = "2026-01-31T23:11:50.293Z" }, + { url = "https://files.pythonhosted.org/packages/4d/89/253db0fa0e66e9129c745e4ef25631dc37d5f1314dad2b53e907b8538e6d/numpy-2.4.2-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:66cb9422236317f9d44b67b4d18f44efe6e9c7f8794ac0462978513359461554", size = 17079064, upload-time = "2026-01-31T23:11:52.927Z" }, + { url = "https://files.pythonhosted.org/packages/2a/d5/cbade46ce97c59c6c3da525e8d95b7abe8a42974a1dc5c1d489c10433e88/numpy-2.4.2-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:0f01dcf33e73d80bd8dc0f20a71303abbafa26a19e23f6b68d1aa9990af90257", size = 18379680, upload-time = "2026-01-31T23:11:55.22Z" }, + { url = "https://files.pythonhosted.org/packages/40/62/48f99ae172a4b63d981babe683685030e8a3df4f246c893ea5c6ef99f018/numpy-2.4.2-cp313-cp313t-win32.whl", hash = "sha256:52b913ec40ff7ae845687b0b34d8d93b60cb66dcee06996dd5c99f2fc9328657", size = 6082433, upload-time = "2026-01-31T23:11:58.096Z" }, + { url = "https://files.pythonhosted.org/packages/07/38/e054a61cfe48ad9f1ed0d188e78b7e26859d0b60ef21cd9de4897cdb5326/numpy-2.4.2-cp313-cp313t-win_amd64.whl", hash = "sha256:5eea80d908b2c1f91486eb95b3fb6fab187e569ec9752ab7d9333d2e66bf2d6b", size = 12451181, upload-time = "2026-01-31T23:11:59.782Z" }, + { url = "https://files.pythonhosted.org/packages/6e/a4/a05c3a6418575e185dd84d0b9680b6bb2e2dc3e4202f036b7b4e22d6e9dc/numpy-2.4.2-cp313-cp313t-win_arm64.whl", hash = "sha256:fd49860271d52127d61197bb50b64f58454e9f578cb4b2c001a6de8b1f50b0b1", size = 10290756, upload-time = "2026-01-31T23:12:02.438Z" }, + { url = "https://files.pythonhosted.org/packages/18/88/b7df6050bf18fdcfb7046286c6535cabbdd2064a3440fca3f069d319c16e/numpy-2.4.2-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:444be170853f1f9d528428eceb55f12918e4fda5d8805480f36a002f1415e09b", size = 16663092, upload-time = "2026-01-31T23:12:04.521Z" }, + { url = "https://files.pythonhosted.org/packages/25/7a/1fee4329abc705a469a4afe6e69b1ef7e915117747886327104a8493a955/numpy-2.4.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:d1240d50adff70c2a88217698ca844723068533f3f5c5fa6ee2e3220e3bdb000", size = 14698770, upload-time = "2026-01-31T23:12:06.96Z" }, + { url = "https://files.pythonhosted.org/packages/fb/0b/f9e49ba6c923678ad5bc38181c08ac5e53b7a5754dbca8e581aa1a56b1ff/numpy-2.4.2-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:7cdde6de52fb6664b00b056341265441192d1291c130e99183ec0d4b110ff8b1", size = 5208562, upload-time = "2026-01-31T23:12:09.632Z" }, + { url = "https://files.pythonhosted.org/packages/7d/12/d7de8f6f53f9bb76997e5e4c069eda2051e3fe134e9181671c4391677bb2/numpy-2.4.2-cp314-cp314-macosx_14_0_x86_64.whl", hash = "sha256:cda077c2e5b780200b6b3e09d0b42205a3d1c68f30c6dceb90401c13bff8fe74", size = 6543710, upload-time = "2026-01-31T23:12:11.969Z" }, + { url = "https://files.pythonhosted.org/packages/09/63/c66418c2e0268a31a4cf8a8b512685748200f8e8e8ec6c507ce14e773529/numpy-2.4.2-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d30291931c915b2ab5717c2974bb95ee891a1cf22ebc16a8006bd59cd210d40a", size = 15677205, upload-time = "2026-01-31T23:12:14.33Z" }, + { url = "https://files.pythonhosted.org/packages/5d/6c/7f237821c9642fb2a04d2f1e88b4295677144ca93285fd76eff3bcba858d/numpy-2.4.2-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bba37bc29d4d85761deed3954a1bc62be7cf462b9510b51d367b769a8c8df325", size = 16611738, upload-time = "2026-01-31T23:12:16.525Z" }, + { url = "https://files.pythonhosted.org/packages/c2/a7/39c4cdda9f019b609b5c473899d87abff092fc908cfe4d1ecb2fcff453b0/numpy-2.4.2-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:b2f0073ed0868db1dcd86e052d37279eef185b9c8db5bf61f30f46adac63c909", size = 17028888, upload-time = "2026-01-31T23:12:19.306Z" }, + { url = "https://files.pythonhosted.org/packages/da/b3/e84bb64bdfea967cc10950d71090ec2d84b49bc691df0025dddb7c26e8e3/numpy-2.4.2-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:7f54844851cdb630ceb623dcec4db3240d1ac13d4990532446761baede94996a", size = 18339556, upload-time = "2026-01-31T23:12:21.816Z" }, + { url = "https://files.pythonhosted.org/packages/88/f5/954a291bc1192a27081706862ac62bb5920fbecfbaa302f64682aa90beed/numpy-2.4.2-cp314-cp314-win32.whl", hash = "sha256:12e26134a0331d8dbd9351620f037ec470b7c75929cb8a1537f6bfe411152a1a", size = 6006899, upload-time = "2026-01-31T23:12:24.14Z" }, + { url = "https://files.pythonhosted.org/packages/05/cb/eff72a91b2efdd1bc98b3b8759f6a1654aa87612fc86e3d87d6fe4f948c4/numpy-2.4.2-cp314-cp314-win_amd64.whl", hash = "sha256:068cdb2d0d644cdb45670810894f6a0600797a69c05f1ac478e8d31670b8ee75", size = 12443072, upload-time = "2026-01-31T23:12:26.33Z" }, + { url = "https://files.pythonhosted.org/packages/37/75/62726948db36a56428fce4ba80a115716dc4fad6a3a4352487f8bb950966/numpy-2.4.2-cp314-cp314-win_arm64.whl", hash = "sha256:6ed0be1ee58eef41231a5c943d7d1375f093142702d5723ca2eb07db9b934b05", size = 10494886, upload-time = "2026-01-31T23:12:28.488Z" }, + { url = "https://files.pythonhosted.org/packages/36/2f/ee93744f1e0661dc267e4b21940870cabfae187c092e1433b77b09b50ac4/numpy-2.4.2-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:98f16a80e917003a12c0580f97b5f875853ebc33e2eaa4bccfc8201ac6869308", size = 14818567, upload-time = "2026-01-31T23:12:30.709Z" }, + { url = "https://files.pythonhosted.org/packages/a7/24/6535212add7d76ff938d8bdc654f53f88d35cddedf807a599e180dcb8e66/numpy-2.4.2-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:20abd069b9cda45874498b245c8015b18ace6de8546bf50dfa8cea1696ed06ef", size = 5328372, upload-time = "2026-01-31T23:12:32.962Z" }, + { url = "https://files.pythonhosted.org/packages/5e/9d/c48f0a035725f925634bf6b8994253b43f2047f6778a54147d7e213bc5a7/numpy-2.4.2-cp314-cp314t-macosx_14_0_x86_64.whl", hash = "sha256:e98c97502435b53741540a5717a6749ac2ada901056c7db951d33e11c885cc7d", size = 6649306, upload-time = "2026-01-31T23:12:34.797Z" }, + { url = "https://files.pythonhosted.org/packages/81/05/7c73a9574cd4a53a25907bad38b59ac83919c0ddc8234ec157f344d57d9a/numpy-2.4.2-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:da6cad4e82cb893db4b69105c604d805e0c3ce11501a55b5e9f9083b47d2ffe8", size = 15722394, upload-time = "2026-01-31T23:12:36.565Z" }, + { url = "https://files.pythonhosted.org/packages/35/fa/4de10089f21fc7d18442c4a767ab156b25c2a6eaf187c0db6d9ecdaeb43f/numpy-2.4.2-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9e4424677ce4b47fe73c8b5556d876571f7c6945d264201180db2dc34f676ab5", size = 16653343, upload-time = "2026-01-31T23:12:39.188Z" }, + { url = "https://files.pythonhosted.org/packages/b8/f9/d33e4ffc857f3763a57aa85650f2e82486832d7492280ac21ba9efda80da/numpy-2.4.2-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:2b8f157c8a6f20eb657e240f8985cc135598b2b46985c5bccbde7616dc9c6b1e", size = 17078045, upload-time = "2026-01-31T23:12:42.041Z" }, + { url = "https://files.pythonhosted.org/packages/c8/b8/54bdb43b6225badbea6389fa038c4ef868c44f5890f95dd530a218706da3/numpy-2.4.2-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:5daf6f3914a733336dab21a05cdec343144600e964d2fcdabaac0c0269874b2a", size = 18380024, upload-time = "2026-01-31T23:12:44.331Z" }, + { url = "https://files.pythonhosted.org/packages/a5/55/6e1a61ded7af8df04016d81b5b02daa59f2ea9252ee0397cb9f631efe9e5/numpy-2.4.2-cp314-cp314t-win32.whl", hash = "sha256:8c50dd1fc8826f5b26a5ee4d77ca55d88a895f4e4819c7ecc2a9f5905047a443", size = 6153937, upload-time = "2026-01-31T23:12:47.229Z" }, + { url = "https://files.pythonhosted.org/packages/45/aa/fa6118d1ed6d776b0983f3ceac9b1a5558e80df9365b1c3aa6d42bf9eee4/numpy-2.4.2-cp314-cp314t-win_amd64.whl", hash = "sha256:fcf92bee92742edd401ba41135185866f7026c502617f422eb432cfeca4fe236", size = 12631844, upload-time = "2026-01-31T23:12:48.997Z" }, + { url = "https://files.pythonhosted.org/packages/32/0a/2ec5deea6dcd158f254a7b372fb09cfba5719419c8d66343bab35237b3fb/numpy-2.4.2-cp314-cp314t-win_arm64.whl", hash = "sha256:1f92f53998a17265194018d1cc321b2e96e900ca52d54c7c77837b71b9465181", size = 10565379, upload-time = "2026-01-31T23:12:51.345Z" }, +] + +[[package]] +name = "nvidia-cublas-cu12" +version = "12.8.4.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/61/e24b560ab2e2eaeb3c839129175fb330dfcfc29e5203196e5541a4c44682/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8ac4e771d5a348c551b2a426eda6193c19aa630236b418086020df5ba9667142", size = 594346921, upload-time = "2025-03-07T01:44:31.254Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.8.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/02/2adcaa145158bf1a8295d83591d22e4103dbfd821bcaf6f3f53151ca4ffa/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ea0cb07ebda26bb9b29ba82cda34849e73c166c18162d3913575b0c9db9a6182", size = 10248621, upload-time = "2025-03-07T01:40:21.213Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.8.93" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/05/6b/32f747947df2da6994e999492ab306a903659555dddc0fbdeb9d71f75e52/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a7756528852ef889772a84c6cd89d41dfa74667e24cca16bb31f8f061e3e9994", size = 88040029, upload-time = "2025-03-07T01:42:13.562Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.8.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0d/9b/a997b638fcd068ad6e4d53b8551a7d30fe8b404d6f1804abf1df69838932/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adade8dcbd0edf427b7204d480d6066d33902cab2a4707dcfc48a2d0fd44ab90", size = 954765, upload-time = "2025-03-07T01:40:01.615Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.10.2.21" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12", marker = "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" }, +] + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.3.3.83" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" }, +] + +[[package]] +name = "nvidia-cufile-cu12" +version = "1.13.1.3" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bb/fe/1bcba1dfbfb8d01be8d93f07bfc502c93fa23afa6fd5ab3fc7c1df71038a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1d069003be650e131b21c932ec3d8969c1715379251f8d23a1860554b1cb24fc", size = 1197834, upload-time = "2025-03-07T01:45:50.723Z" }, +] + +[[package]] +name = "nvidia-curand-cu12" +version = "10.3.9.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/aa/6584b56dc84ebe9cf93226a5cde4d99080c8e90ab40f0c27bda7a0f29aa1/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b32331d4f4df5d6eefa0554c565b626c7216f87a06a4f56fab27c3b68a830ec9", size = 63619976, upload-time = "2025-03-07T01:46:23.323Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.7.3.90" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12", marker = "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'" }, + { name = "nvidia-cusparse-cu12", marker = "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'" }, + { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" }, +] + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.5.8.93" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12", marker = "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" }, +] + +[[package]] +name = "nvidia-cusparselt-cu12" +version = "0.7.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691, upload-time = "2025-02-26T00:15:44.104Z" }, +] + +[[package]] +name = "nvidia-nccl-cu12" +version = "2.27.5" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6e/89/f7a07dc961b60645dbbf42e80f2bc85ade7feb9a491b11a1e973aa00071f/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ad730cf15cb5d25fe849c6e6ca9eb5b76db16a80f13f425ac68d8e2e55624457", size = 322348229, upload-time = "2025-06-26T04:11:28.385Z" }, +] + +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.8.93" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f6/74/86a07f1d0f42998ca31312f998bd3b9a7eff7f52378f4f270c8679c77fb9/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:81ff63371a7ebd6e6451970684f916be2eab07321b73c9d244dc2b4da7f73b88", size = 39254836, upload-time = "2025-03-07T01:49:55.661Z" }, +] + +[[package]] +name = "nvidia-nvshmem-cu12" +version = "3.3.20" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3b/6c/99acb2f9eb85c29fc6f3a7ac4dccfd992e22666dd08a642b303311326a97/nvidia_nvshmem_cu12-3.3.20-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:d00f26d3f9b2e3c3065be895e3059d6479ea5c638a3f38c9fec49b1b9dd7c1e5", size = 124657145, upload-time = "2025-08-04T20:25:19.995Z" }, +] + +[[package]] +name = "nvidia-nvtx-cu12" +version = "12.8.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" }, +] + +[[package]] +name = "packaging" +version = "26.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/65/ee/299d360cdc32edc7d2cf530f3accf79c4fca01e96ffc950d8a52213bd8e4/packaging-26.0.tar.gz", hash = "sha256:00243ae351a257117b6a241061796684b084ed1c516a08c48a3f7e147a9d80b4", size = 143416, upload-time = "2026-01-21T20:50:39.064Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/b9/c538f279a4e237a006a2c98387d081e9eb060d203d8ed34467cc0f0b9b53/packaging-26.0-py3-none-any.whl", hash = "sha256:b36f1fef9334a5588b4166f8bcd26a14e521f2b55e6b9de3aaa80d3ff7a37529", size = 74366, upload-time = "2026-01-21T20:50:37.788Z" }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + +[[package]] +name = "pygame" +version = "2.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/49/cc/08bba60f00541f62aaa252ce0cfbd60aebd04616c0b9574f755b583e45ae/pygame-2.6.1.tar.gz", hash = "sha256:56fb02ead529cee00d415c3e007f75e0780c655909aaa8e8bf616ee09c9feb1f", size = 14808125, upload-time = "2024-09-29T13:41:34.698Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/92/16/2c602c332f45ff9526d61f6bd764db5096ff9035433e2172e2d2cadae8db/pygame-2.6.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:4ee7f2771f588c966fa2fa8b829be26698c9b4836f82ede5e4edc1a68594942e", size = 13118279, upload-time = "2024-09-29T14:26:30.427Z" }, + { url = "https://files.pythonhosted.org/packages/cd/53/77ccbc384b251c6e34bfd2e734c638233922449a7844e3c7a11ef91cee39/pygame-2.6.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c8040ea2ab18c6b255af706ec01355c8a6b08dc48d77fd4ee783f8fc46a843bf", size = 12384524, upload-time = "2024-09-29T14:26:49.996Z" }, + { url = "https://files.pythonhosted.org/packages/06/be/3ed337583f010696c3b3435e89a74fb29d0c74d0931e8f33c0a4246307a9/pygame-2.6.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c47a6938de93fa610accd4969e638c2aebcb29b2fca518a84c3a39d91ab47116", size = 13587123, upload-time = "2024-09-29T11:10:50.072Z" }, + { url = "https://files.pythonhosted.org/packages/fd/ca/b015586a450db59313535662991b34d24c1f0c0dc149cc5f496573900f4e/pygame-2.6.1-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:33006f784e1c7d7e466fcb61d5489da59cc5f7eb098712f792a225df1d4e229d", size = 14275532, upload-time = "2024-09-29T11:39:59.356Z" }, + { url = "https://files.pythonhosted.org/packages/b9/f2/d31e6ad42d657af07be2ffd779190353f759a07b51232b9e1d724f2cda46/pygame-2.6.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1206125f14cae22c44565c9d333607f1d9f59487b1f1432945dfc809aeaa3e88", size = 13952653, upload-time = "2024-09-29T11:40:01.781Z" }, + { url = "https://files.pythonhosted.org/packages/f3/42/8ea2a6979e6fa971702fece1747e862e2256d4a8558fe0da6364dd946c53/pygame-2.6.1-cp312-cp312-win32.whl", hash = "sha256:84fc4054e25262140d09d39e094f6880d730199710829902f0d8ceae0213379e", size = 10252421, upload-time = "2024-09-29T11:14:26.877Z" }, + { url = "https://files.pythonhosted.org/packages/5f/90/7d766d54bb95939725e9a9361f9c06b0cfbe3fe100aa35400f0a461a278a/pygame-2.6.1-cp312-cp312-win_amd64.whl", hash = "sha256:3a9e7396be0d9633831c3f8d5d82dd63ba373ad65599628294b7a4f8a5a01a65", size = 10624591, upload-time = "2024-09-29T11:52:54.489Z" }, + { url = "https://files.pythonhosted.org/packages/e1/91/718acf3e2a9d08a6ddcc96bd02a6f63c99ee7ba14afeaff2a51c987df0b9/pygame-2.6.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:ae6039f3a55d800db80e8010f387557b528d34d534435e0871326804df2a62f2", size = 13090765, upload-time = "2024-09-29T14:27:02.377Z" }, + { url = "https://files.pythonhosted.org/packages/0e/c6/9cb315de851a7682d9c7568a41ea042ee98d668cb8deadc1dafcab6116f0/pygame-2.6.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:2a3a1288e2e9b1e5834e425bedd5ba01a3cd4902b5c2bff8ed4a740ccfe98171", size = 12381704, upload-time = "2024-09-29T14:27:10.228Z" }, + { url = "https://files.pythonhosted.org/packages/9f/8f/617a1196e31ae3b46be6949fbaa95b8c93ce15e0544266198c2266cc1b4d/pygame-2.6.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:27eb17e3dc9640e4b4683074f1890e2e879827447770470c2aba9f125f74510b", size = 13581091, upload-time = "2024-09-29T11:30:27.653Z" }, + { url = "https://files.pythonhosted.org/packages/3b/87/2851a564e40a2dad353f1c6e143465d445dab18a95281f9ea458b94f3608/pygame-2.6.1-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4c1623180e70a03c4a734deb9bac50fc9c82942ae84a3a220779062128e75f3b", size = 14273844, upload-time = "2024-09-29T11:40:04.138Z" }, + { url = "https://files.pythonhosted.org/packages/85/b5/aa23aa2e70bcba42c989c02e7228273c30f3b44b9b264abb93eaeff43ad7/pygame-2.6.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ef07c0103d79492c21fced9ad68c11c32efa6801ca1920ebfd0f15fb46c78b1c", size = 13951197, upload-time = "2024-09-29T11:40:06.785Z" }, + { url = "https://files.pythonhosted.org/packages/a6/06/29e939b34d3f1354738c7d201c51c250ad7abefefaf6f8332d962ff67c4b/pygame-2.6.1-cp313-cp313-win32.whl", hash = "sha256:3acd8c009317190c2bfd81db681ecef47d5eb108c2151d09596d9c7ea9df5c0e", size = 10249309, upload-time = "2024-09-29T11:10:23.329Z" }, + { url = "https://files.pythonhosted.org/packages/7e/11/17f7f319ca91824b86557e9303e3b7a71991ef17fd45286bf47d7f0a38e6/pygame-2.6.1-cp313-cp313-win_amd64.whl", hash = "sha256:813af4fba5d0b2cb8e58f5d95f7910295c34067dcc290d34f1be59c48bd1ea6a", size = 10620084, upload-time = "2024-09-29T11:48:51.587Z" }, +] + +[[package]] +name = "pygments" +version = "2.19.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b0/77/a5b8c569bf593b0140bde72ea885a803b82086995367bf2037de0159d924/pygments-2.19.2.tar.gz", hash = "sha256:636cb2477cec7f8952536970bc533bc43743542f70392ae026374600add5b887", size = 4968631, upload-time = "2025-06-21T13:39:12.283Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c7/21/705964c7812476f378728bdf590ca4b771ec72385c533964653c68e86bdc/pygments-2.19.2-py3-none-any.whl", hash = "sha256:86540386c03d588bb81d44bc3928634ff26449851e99741617ecb9037ee5ec0b", size = 1225217, upload-time = "2025-06-21T13:39:07.939Z" }, +] + +[[package]] +name = "pytest" +version = "9.0.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "iniconfig" }, + { name = "packaging" }, + { name = "pluggy" }, + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d1/db/7ef3487e0fb0049ddb5ce41d3a49c235bf9ad299b6a25d5780a89f19230f/pytest-9.0.2.tar.gz", hash = "sha256:75186651a92bd89611d1d9fc20f0b4345fd827c41ccd5c299a868a05d70edf11", size = 1568901, upload-time = "2025-12-06T21:30:51.014Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3b/ab/b3226f0bd7cdcf710fbede2b3548584366da3b19b5021e74f5bde2a8fa3f/pytest-9.0.2-py3-none-any.whl", hash = "sha256:711ffd45bf766d5264d487b917733b453d917afd2b0ad65223959f59089f875b", size = 374801, upload-time = "2025-12-06T21:30:49.154Z" }, +] + +[[package]] +name = "python-chess" +version = "1.999" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "chess" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5f/60/7c7d132b6683ff215bf705fc55ffc0240a6cddea89657407ca0a4fb628d0/python-chess-1.999.tar.gz", hash = "sha256:8cad0388c42242d890ac6368ad64def15cd0165db033df0ad479492e266e5e6c", size = 1453, upload-time = "2020-10-26T11:30:10.118Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/38/47/dfebc06e589530691d33c71bc7b3d8d311252b58ae338980fd4327fe77f2/python_chess-1.999-py3-none-any.whl", hash = "sha256:93b562f8f1124cb7bf56fb095e18743758e69dc6a028ccda0badcaa5c59d88c8", size = 1401, upload-time = "2020-10-26T11:30:07.758Z" }, +] + +[[package]] +name = "setuptools" +version = "82.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/82/f3/748f4d6f65d1756b9ae577f329c951cda23fb900e4de9f70900ced962085/setuptools-82.0.0.tar.gz", hash = "sha256:22e0a2d69474c6ae4feb01951cb69d515ed23728cf96d05513d36e42b62b37cb", size = 1144893, upload-time = "2026-02-08T15:08:40.206Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e1/c6/76dc613121b793286a3f91621d7b75a2b493e0390ddca50f11993eadf192/setuptools-82.0.0-py3-none-any.whl", hash = "sha256:70b18734b607bd1da571d097d236cfcfacaf01de45717d59e6e04b96877532e0", size = 1003468, upload-time = "2026-02-08T15:08:38.723Z" }, +] + +[[package]] +name = "sympy" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mpmath" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, +] + +[[package]] +name = "torch" +version = "2.9.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'", +] +dependencies = [ + { name = "filelock", marker = "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'" }, + { name = "fsspec", marker = "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'" }, + { name = "jinja2", marker = "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'" }, + { name = "networkx", marker = "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'" }, + { name = "nvidia-cublas-cu12", marker = "python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", marker = "python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", marker = "python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "nvidia-cufile-cu12", marker = "python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "nvidia-curand-cu12", marker = "python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", marker = "python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu12", marker = "python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", marker = "python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu12", marker = "python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "nvidia-nvtx-cu12", marker = "python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "setuptools", marker = "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'" }, + { name = "sympy", marker = "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'" }, + { name = "triton", version = "3.5.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux'" }, + { name = "typing-extensions", marker = "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/0f/27/07c645c7673e73e53ded71705045d6cb5bae94c4b021b03aa8d03eee90ab/torch-2.9.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:da5f6f4d7f4940a173e5572791af238cb0b9e21b1aab592bd8b26da4c99f1cd6", size = 104126592, upload-time = "2025-11-12T15:20:41.62Z" }, + { url = "https://files.pythonhosted.org/packages/19/17/e377a460603132b00760511299fceba4102bd95db1a0ee788da21298ccff/torch-2.9.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:27331cd902fb4322252657f3902adf1c4f6acad9dcad81d8df3ae14c7c4f07c4", size = 899742281, upload-time = "2025-11-12T15:22:17.602Z" }, + { url = "https://files.pythonhosted.org/packages/b1/1a/64f5769025db846a82567fa5b7d21dba4558a7234ee631712ee4771c436c/torch-2.9.1-cp312-cp312-win_amd64.whl", hash = "sha256:81a285002d7b8cfd3fdf1b98aa8df138d41f1a8334fd9ea37511517cedf43083", size = 110940568, upload-time = "2025-11-12T15:21:18.689Z" }, + { url = "https://files.pythonhosted.org/packages/6e/ab/07739fd776618e5882661d04c43f5b5586323e2f6a2d7d84aac20d8f20bd/torch-2.9.1-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:c0d25d1d8e531b8343bea0ed811d5d528958f1dcbd37e7245bc686273177ad7e", size = 74479191, upload-time = "2025-11-12T15:21:25.816Z" }, + { url = "https://files.pythonhosted.org/packages/20/60/8fc5e828d050bddfab469b3fe78e5ab9a7e53dda9c3bdc6a43d17ce99e63/torch-2.9.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:c29455d2b910b98738131990394da3e50eea8291dfeb4b12de71ecf1fdeb21cb", size = 104135743, upload-time = "2025-11-12T15:21:34.936Z" }, + { url = "https://files.pythonhosted.org/packages/f2/b7/6d3f80e6918213babddb2a37b46dbb14c15b14c5f473e347869a51f40e1f/torch-2.9.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:524de44cd13931208ba2c4bde9ec7741fd4ae6bfd06409a604fc32f6520c2bc9", size = 899749493, upload-time = "2025-11-12T15:24:36.356Z" }, + { url = "https://files.pythonhosted.org/packages/a6/47/c7843d69d6de8938c1cbb1eba426b1d48ddf375f101473d3e31a5fc52b74/torch-2.9.1-cp313-cp313-win_amd64.whl", hash = "sha256:545844cc16b3f91e08ce3b40e9c2d77012dd33a48d505aed34b7740ed627a1b2", size = 110944162, upload-time = "2025-11-12T15:21:53.151Z" }, + { url = "https://files.pythonhosted.org/packages/28/0e/2a37247957e72c12151b33a01e4df651d9d155dd74d8cfcbfad15a79b44a/torch-2.9.1-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:5be4bf7496f1e3ffb1dd44b672adb1ac3f081f204c5ca81eba6442f5f634df8e", size = 74830751, upload-time = "2025-11-12T15:21:43.792Z" }, + { url = "https://files.pythonhosted.org/packages/4b/f7/7a18745edcd7b9ca2381aa03353647bca8aace91683c4975f19ac233809d/torch-2.9.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:30a3e170a84894f3652434b56d59a64a2c11366b0ed5776fab33c2439396bf9a", size = 104142929, upload-time = "2025-11-12T15:21:48.319Z" }, + { url = "https://files.pythonhosted.org/packages/f4/dd/f1c0d879f2863ef209e18823a988dc7a1bf40470750e3ebe927efdb9407f/torch-2.9.1-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:8301a7b431e51764629208d0edaa4f9e4c33e6df0f2f90b90e261d623df6a4e2", size = 899748978, upload-time = "2025-11-12T15:23:04.568Z" }, + { url = "https://files.pythonhosted.org/packages/1f/9f/6986b83a53b4d043e36f3f898b798ab51f7f20fdf1a9b01a2720f445043d/torch-2.9.1-cp313-cp313t-win_amd64.whl", hash = "sha256:2e1c42c0ae92bf803a4b2409fdfed85e30f9027a66887f5e7dcdbc014c7531db", size = 111176995, upload-time = "2025-11-12T15:22:01.618Z" }, + { url = "https://files.pythonhosted.org/packages/40/60/71c698b466dd01e65d0e9514b5405faae200c52a76901baf6906856f17e4/torch-2.9.1-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:2c14b3da5df416cf9cb5efab83aa3056f5b8cd8620b8fde81b4987ecab730587", size = 74480347, upload-time = "2025-11-12T15:21:57.648Z" }, + { url = "https://files.pythonhosted.org/packages/48/50/c4b5112546d0d13cc9eaa1c732b823d676a9f49ae8b6f97772f795874a03/torch-2.9.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:1edee27a7c9897f4e0b7c14cfc2f3008c571921134522d5b9b5ec4ebbc69041a", size = 74433245, upload-time = "2025-11-12T15:22:39.027Z" }, + { url = "https://files.pythonhosted.org/packages/81/c9/2628f408f0518b3bae49c95f5af3728b6ab498c8624ab1e03a43dd53d650/torch-2.9.1-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:19d144d6b3e29921f1fc70503e9f2fc572cde6a5115c0c0de2f7ca8b1483e8b6", size = 104134804, upload-time = "2025-11-12T15:22:35.222Z" }, + { url = "https://files.pythonhosted.org/packages/28/fc/5bc91d6d831ae41bf6e9e6da6468f25330522e92347c9156eb3f1cb95956/torch-2.9.1-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:c432d04376f6d9767a9852ea0def7b47a7bbc8e7af3b16ac9cf9ce02b12851c9", size = 899747132, upload-time = "2025-11-12T15:23:36.068Z" }, + { url = "https://files.pythonhosted.org/packages/63/5d/e8d4e009e52b6b2cf1684bde2a6be157b96fb873732542fb2a9a99e85a83/torch-2.9.1-cp314-cp314-win_amd64.whl", hash = "sha256:d187566a2cdc726fc80138c3cdb260970fab1c27e99f85452721f7759bbd554d", size = 110934845, upload-time = "2025-11-12T15:22:48.367Z" }, + { url = "https://files.pythonhosted.org/packages/bd/b2/2d15a52516b2ea3f414643b8de68fa4cb220d3877ac8b1028c83dc8ca1c4/torch-2.9.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:cb10896a1f7fedaddbccc2017ce6ca9ecaaf990f0973bdfcf405439750118d2c", size = 74823558, upload-time = "2025-11-12T15:22:43.392Z" }, + { url = "https://files.pythonhosted.org/packages/86/5c/5b2e5d84f5b9850cd1e71af07524d8cbb74cba19379800f1f9f7c997fc70/torch-2.9.1-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:0a2bd769944991c74acf0c4ef23603b9c777fdf7637f115605a4b2d8023110c7", size = 104145788, upload-time = "2025-11-12T15:23:52.109Z" }, + { url = "https://files.pythonhosted.org/packages/a9/8c/3da60787bcf70add986c4ad485993026ac0ca74f2fc21410bc4eb1bb7695/torch-2.9.1-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:07c8a9660bc9414c39cac530ac83b1fb1b679d7155824144a40a54f4a47bfa73", size = 899735500, upload-time = "2025-11-12T15:24:08.788Z" }, + { url = "https://files.pythonhosted.org/packages/db/2b/f7818f6ec88758dfd21da46b6cd46af9d1b3433e53ddbb19ad1e0da17f9b/torch-2.9.1-cp314-cp314t-win_amd64.whl", hash = "sha256:c88d3299ddeb2b35dcc31753305612db485ab6f1823e37fb29451c8b2732b87e", size = 111163659, upload-time = "2025-11-12T15:23:20.009Z" }, +] + +[[package]] +name = "torch" +version = "2.9.1+rocm7.2.0.lw.git7e1940d4" +source = { registry = "wheelies" } +resolution-markers = [ + "(python_full_version >= '3.15' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')", +] +dependencies = [ + { name = "filelock", marker = "(python_full_version >= '3.15' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')" }, + { name = "fsspec", marker = "(python_full_version >= '3.15' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')" }, + { name = "jinja2", marker = "(python_full_version >= '3.15' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')" }, + { name = "networkx", marker = "(python_full_version >= '3.15' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')" }, + { name = "setuptools", marker = "(python_full_version >= '3.15' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')" }, + { name = "sympy", marker = "(python_full_version >= '3.15' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')" }, + { name = "triton", version = "3.5.1+rocm7.2.0.gita272dfa8", source = { registry = "wheelies" }, marker = "(python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'CPython' and sys_platform == 'linux')" }, + { name = "typing-extensions", marker = "(python_full_version >= '3.15' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')" }, +] +wheels = [ + { path = "torch-2.9.1+rocm7.2.0.lw.git7e1940d4-cp312-cp312-linux_x86_64.whl" }, +] + +[[package]] +name = "triton" +version = "3.5.1" +source = { registry = "https://pypi.org/simple" } +resolution-markers = [ + "(python_full_version >= '3.13' and python_full_version < '3.15' and platform_machine == 'x86_64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or (python_full_version < '3.15' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux') or sys_platform == 'darwin' or sys_platform == 'win32'", +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/f2/50/9a8358d3ef58162c0a415d173cfb45b67de60176e1024f71fbc4d24c0b6d/triton-3.5.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d2c6b915a03888ab931a9fd3e55ba36785e1fe70cbea0b40c6ef93b20fc85232", size = 170470207, upload-time = "2025-11-11T17:41:00.253Z" }, + { url = "https://files.pythonhosted.org/packages/27/46/8c3bbb5b0a19313f50edcaa363b599e5a1a5ac9683ead82b9b80fe497c8d/triton-3.5.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f3f4346b6ebbd4fad18773f5ba839114f4826037c9f2f34e0148894cd5dd3dba", size = 170470410, upload-time = "2025-11-11T17:41:06.319Z" }, + { url = "https://files.pythonhosted.org/packages/37/92/e97fcc6b2c27cdb87ce5ee063d77f8f26f19f06916aa680464c8104ef0f6/triton-3.5.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0b4d2c70127fca6a23e247f9348b8adde979d2e7a20391bfbabaac6aebc7e6a8", size = 170579924, upload-time = "2025-11-11T17:41:12.455Z" }, + { url = "https://files.pythonhosted.org/packages/a4/e6/c595c35e5c50c4bc56a7bac96493dad321e9e29b953b526bbbe20f9911d0/triton-3.5.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d0637b1efb1db599a8e9dc960d53ab6e4637db7d4ab6630a0974705d77b14b60", size = 170480488, upload-time = "2025-11-11T17:41:18.222Z" }, + { url = "https://files.pythonhosted.org/packages/16/b5/b0d3d8b901b6a04ca38df5e24c27e53afb15b93624d7fd7d658c7cd9352a/triton-3.5.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bac7f7d959ad0f48c0e97d6643a1cc0fd5786fe61cb1f83b537c6b2d54776478", size = 170582192, upload-time = "2025-11-11T17:41:23.963Z" }, +] + +[[package]] +name = "triton" +version = "3.5.1+rocm7.2.0.gita272dfa8" +source = { registry = "wheelies" } +resolution-markers = [ + "(python_full_version >= '3.15' and platform_machine == 'aarch64' and sys_platform == 'linux') or (python_full_version < '3.13' and platform_machine == 'x86_64' and sys_platform == 'linux') or (python_full_version >= '3.15' and platform_machine == 'x86_64' and sys_platform == 'linux') or (platform_machine != 'aarch64' and platform_machine != 'x86_64' and sys_platform == 'linux') or (platform_machine == 'aarch64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (platform_machine == 'x86_64' and platform_python_implementation != 'CPython' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux' and sys_platform != 'win32')", +] +wheels = [ + { path = "triton-3.5.1+rocm7.2.0.gita272dfa8-cp312-cp312-linux_x86_64.whl" }, +] + +[[package]] +name = "typing-extensions" +version = "4.15.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/94/1a15dd82efb362ac84269196e94cf00f187f7ed21c242792a923cdb1c61f/typing_extensions-4.15.0.tar.gz", hash = "sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466", size = 109391, upload-time = "2025-08-25T13:49:26.313Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" }, +]