276 lines
8.6 KiB
Python
276 lines
8.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Central NNUE pipeline CLI for training and exporting models."""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import subprocess
|
|
from pathlib import Path
|
|
|
|
def get_python_cmd():
|
|
"""Get available Python command."""
|
|
if os.name == 'nt':
|
|
return "python"
|
|
return "python3" if os.popen("which python3 2>/dev/null").read() else "python"
|
|
|
|
def get_src_module(module_name):
|
|
"""Get path to module in src/ directory."""
|
|
return Path(__file__).parent / "src" / f"{module_name}.py"
|
|
|
|
def get_data_dir():
|
|
"""Get/create data directory."""
|
|
data_dir = Path(__file__).parent / "data"
|
|
data_dir.mkdir(exist_ok=True)
|
|
return data_dir
|
|
|
|
def get_weights_dir():
|
|
"""Get/create weights directory."""
|
|
weights_dir = Path(__file__).parent / "weights"
|
|
weights_dir.mkdir(exist_ok=True)
|
|
return weights_dir
|
|
|
|
def list_checkpoints():
|
|
"""List available checkpoint versions."""
|
|
weights_dir = get_weights_dir()
|
|
checkpoints = sorted(weights_dir.glob("nnue_weights_v*.pt"))
|
|
if not checkpoints:
|
|
return []
|
|
return [int(cp.stem.split("_v")[1]) for cp in checkpoints]
|
|
|
|
def run_generate_positions(num_games):
|
|
"""Generate random positions."""
|
|
data_dir = get_data_dir()
|
|
positions_file = data_dir / "positions.txt"
|
|
print(f"Generating {num_games} positions...")
|
|
result = subprocess.run(
|
|
[get_python_cmd(), str(get_src_module("generate")), str(positions_file), "--games", str(num_games)],
|
|
capture_output=False
|
|
)
|
|
if result.returncode != 0:
|
|
print("ERROR: Position generation failed")
|
|
return False
|
|
return positions_file.exists()
|
|
|
|
def run_label_positions(stockfish_path):
|
|
"""Label positions with Stockfish."""
|
|
data_dir = get_data_dir()
|
|
positions_file = data_dir / "positions.txt"
|
|
output_file = data_dir / "training_data.jsonl"
|
|
|
|
if not positions_file.exists():
|
|
print("ERROR: positions.txt not found")
|
|
return False
|
|
|
|
print("Labeling positions with Stockfish...")
|
|
result = subprocess.run(
|
|
[get_python_cmd(), str(get_src_module("label")), str(positions_file), str(output_file), stockfish_path],
|
|
capture_output=False
|
|
)
|
|
if result.returncode != 0:
|
|
print("ERROR: Position labeling failed")
|
|
return False
|
|
return output_file.exists()
|
|
|
|
def run_train(positions_file, output_weights, from_checkpoint=None):
|
|
"""Train NNUE model."""
|
|
if not Path(positions_file).exists():
|
|
print(f"ERROR: {positions_file} not found")
|
|
return False
|
|
|
|
weights_dir = get_weights_dir()
|
|
print(f"Training model (output: {output_weights})...")
|
|
if from_checkpoint:
|
|
print(f" Starting from checkpoint: {from_checkpoint}")
|
|
|
|
cmd = [get_python_cmd(), str(get_src_module("train")), str(positions_file), str(output_weights)]
|
|
if from_checkpoint:
|
|
cmd.extend(["--checkpoint", str(from_checkpoint)])
|
|
|
|
# Run from weights directory so outputs save there
|
|
result = subprocess.run(cmd, cwd=str(weights_dir), capture_output=False)
|
|
if result.returncode != 0:
|
|
print("ERROR: Training failed")
|
|
return False
|
|
return True
|
|
|
|
def run_export(weights_file, output_file):
|
|
"""Export weights to Scala."""
|
|
weights_dir = get_weights_dir()
|
|
weights_path = weights_dir / Path(weights_file).name
|
|
|
|
if not weights_path.exists():
|
|
print(f"ERROR: {weights_file} not found in {weights_dir}")
|
|
return False
|
|
|
|
print(f"Exporting {weights_file} to Scala...")
|
|
result = subprocess.run(
|
|
[get_python_cmd(), str(get_src_module("export")), str(weights_path), output_file],
|
|
capture_output=False
|
|
)
|
|
if result.returncode != 0:
|
|
print("ERROR: Export failed")
|
|
return False
|
|
return Path(output_file).exists()
|
|
|
|
def cmd_train(args):
|
|
"""Handle train command."""
|
|
stockfish_path = args.stockfish or os.environ.get("STOCKFISH_PATH", "/usr/games/stockfish")
|
|
data_dir = get_data_dir()
|
|
weights_dir = get_weights_dir()
|
|
|
|
# Determine checkpoint
|
|
checkpoint = None
|
|
if args.from_checkpoint:
|
|
checkpoint_version = args.from_checkpoint
|
|
checkpoint = f"nnue_weights_v{checkpoint_version}.pt"
|
|
checkpoint_path = weights_dir / checkpoint
|
|
if not checkpoint_path.exists():
|
|
print(f"ERROR: Checkpoint {checkpoint} not found")
|
|
return False
|
|
else:
|
|
available = list_checkpoints()
|
|
if available:
|
|
latest = max(available)
|
|
checkpoint = f"nnue_weights_v{latest}.pt"
|
|
print(f"No checkpoint specified, using latest: v{latest}")
|
|
|
|
# Generate or use existing positions
|
|
if args.positions_file:
|
|
positions_file = Path(args.positions_file)
|
|
if not positions_file.exists():
|
|
print(f"ERROR: {args.positions_file} not found")
|
|
return False
|
|
else:
|
|
positions_file = data_dir / "positions.txt"
|
|
num_games = args.games or 500000
|
|
if not run_generate_positions(num_games):
|
|
return False
|
|
|
|
# Label positions
|
|
if not run_label_positions(stockfish_path):
|
|
return False
|
|
|
|
print("\nStarting training...")
|
|
|
|
# Train with absolute path to data, checkpoint is relative to weights dir
|
|
training_data = str(data_dir / "training_data.jsonl")
|
|
if not run_train(training_data, "nnue_weights.pt", checkpoint):
|
|
return False
|
|
|
|
# Show created version
|
|
available = list_checkpoints()
|
|
new_version = max(available) if available else 1
|
|
print(f"\n✓ Training complete: nnue_weights_v{new_version}.pt")
|
|
return True
|
|
|
|
def cmd_export(args):
|
|
"""Handle export command."""
|
|
weights_file = args.weights
|
|
|
|
# Auto-detect if version is specified
|
|
if not weights_file.endswith(".pt"):
|
|
weights_file = f"nnue_weights_v{weights_file}.pt"
|
|
|
|
# Output to resources directory as binary format
|
|
output_file = str(Path(__file__).parent.parent / "src" / "main" / "resources" / "nnue_weights.bin")
|
|
|
|
if not run_export(weights_file, output_file):
|
|
return False
|
|
|
|
print(f"✓ Export complete: {output_file}")
|
|
return True
|
|
|
|
def cmd_list(args):
|
|
"""List available checkpoints."""
|
|
available = list_checkpoints()
|
|
if not available:
|
|
print("No checkpoints found")
|
|
return True
|
|
|
|
weights_dir = get_weights_dir()
|
|
print("Available checkpoints:")
|
|
for v in available:
|
|
weights_file = weights_dir / f"nnue_weights_v{v}.pt"
|
|
if weights_file.exists():
|
|
size = weights_file.stat().st_size / (1024**2) # MB
|
|
print(f" v{v} ({size:.1f} MB)")
|
|
else:
|
|
print(f" v{v} (file not found)")
|
|
return True
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="NNUE pipeline CLI for training and exporting models",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog="""
|
|
Examples:
|
|
# Train with 500k random positions
|
|
python nnue.py train
|
|
|
|
# Train from checkpoint v2
|
|
python nnue.py train --from-checkpoint 2
|
|
|
|
# Train with custom positions file
|
|
python nnue.py train --positions-file my_positions.txt
|
|
|
|
# Train with 200k games
|
|
python nnue.py train --games 200000
|
|
|
|
# Export specific weights version
|
|
python nnue.py export 2
|
|
|
|
# Export with full filename
|
|
python nnue.py export nnue_weights_v3.pt
|
|
|
|
# List available checkpoints
|
|
python nnue.py list
|
|
"""
|
|
)
|
|
|
|
subparsers = parser.add_subparsers(dest="command", help="Command to run")
|
|
|
|
# Train subcommand
|
|
train_parser = subparsers.add_parser("train", help="Train NNUE model")
|
|
train_parser.add_argument(
|
|
"--from-checkpoint",
|
|
type=int,
|
|
help="Start training from checkpoint version (e.g., 2)"
|
|
)
|
|
train_parser.add_argument(
|
|
"--games",
|
|
type=int,
|
|
help="Number of games to generate (default: 500000)"
|
|
)
|
|
train_parser.add_argument(
|
|
"--positions-file",
|
|
help="Use existing positions file instead of generating"
|
|
)
|
|
train_parser.add_argument(
|
|
"--stockfish",
|
|
help="Path to Stockfish binary (default: $STOCKFISH_PATH or /usr/games/stockfish)"
|
|
)
|
|
train_parser.set_defaults(func=cmd_train)
|
|
|
|
# Export subcommand
|
|
export_parser = subparsers.add_parser("export", help="Export weights to Scala")
|
|
export_parser.add_argument(
|
|
"weights",
|
|
help="Weights file or version (e.g., 2 or nnue_weights_v2.pt)"
|
|
)
|
|
export_parser.set_defaults(func=cmd_export)
|
|
|
|
# List subcommand
|
|
list_parser = subparsers.add_parser("list", help="List available checkpoints")
|
|
list_parser.set_defaults(func=cmd_list)
|
|
|
|
args = parser.parse_args()
|
|
|
|
if not args.command:
|
|
parser.print_help()
|
|
return 0
|
|
|
|
success = args.func(args)
|
|
return 0 if success else 1
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main()) |