feat: refactor AlphaBetaSearch and ClassicalBot for improved evaluation and organization
This commit is contained in:
@@ -0,0 +1,249 @@
|
||||
#!/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 list_checkpoints():
|
||||
"""List available checkpoint versions."""
|
||||
checkpoints = sorted(Path(".").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."""
|
||||
positions_file = "positions.txt"
|
||||
print(f"Generating {num_games} positions...")
|
||||
result = subprocess.run(
|
||||
[get_python_cmd(), "generate_positions.py", positions_file, "--games", str(num_games)],
|
||||
capture_output=False
|
||||
)
|
||||
if result.returncode != 0:
|
||||
print("ERROR: Position generation failed")
|
||||
return False
|
||||
return Path(positions_file).exists()
|
||||
|
||||
def run_label_positions(stockfish_path):
|
||||
"""Label positions with Stockfish."""
|
||||
positions_file = "positions.txt"
|
||||
output_file = "training_data.jsonl"
|
||||
|
||||
if not Path(positions_file).exists():
|
||||
print("ERROR: positions.txt not found")
|
||||
return False
|
||||
|
||||
print("Labeling positions with Stockfish...")
|
||||
result = subprocess.run(
|
||||
[get_python_cmd(), "label_positions.py", positions_file, output_file, stockfish_path],
|
||||
capture_output=False
|
||||
)
|
||||
if result.returncode != 0:
|
||||
print("ERROR: Position labeling failed")
|
||||
return False
|
||||
return Path(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
|
||||
|
||||
print(f"Training model (output: {output_weights})...")
|
||||
if from_checkpoint:
|
||||
print(f" Starting from checkpoint: {from_checkpoint}")
|
||||
|
||||
cmd = [get_python_cmd(), "train_nnue.py", positions_file, output_weights]
|
||||
if from_checkpoint:
|
||||
cmd.extend(["--checkpoint", from_checkpoint])
|
||||
|
||||
result = subprocess.run(cmd, capture_output=False)
|
||||
if result.returncode != 0:
|
||||
print("ERROR: Training failed")
|
||||
return False
|
||||
return True # train_nnue creates versioned file, not the base name
|
||||
|
||||
def run_export(weights_file, output_file):
|
||||
"""Export weights to Scala."""
|
||||
if not Path(weights_file).exists():
|
||||
print(f"ERROR: {weights_file} not found")
|
||||
return False
|
||||
|
||||
print(f"Exporting {weights_file} to Scala...")
|
||||
result = subprocess.run(
|
||||
[get_python_cmd(), "export_weights.py", weights_file, 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")
|
||||
|
||||
# Determine checkpoint
|
||||
checkpoint = None
|
||||
if args.from_checkpoint:
|
||||
checkpoint_version = args.from_checkpoint
|
||||
checkpoint = f"nnue_weights_v{checkpoint_version}.pt"
|
||||
if not Path(checkpoint).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:
|
||||
if not Path(args.positions_file).exists():
|
||||
print(f"ERROR: {args.positions_file} not found")
|
||||
return False
|
||||
positions_file = args.positions_file
|
||||
else:
|
||||
positions_file = "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 (train_nnue.py handles versioning internally)
|
||||
if not run_train("training_data.jsonl", "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"
|
||||
|
||||
if not Path(weights_file).exists():
|
||||
print(f"ERROR: {weights_file} not found")
|
||||
return False
|
||||
|
||||
# Determine version from filename
|
||||
version = Path(weights_file).stem.split("_v")[1] if "_v" in weights_file else "1"
|
||||
output_file = f"../src/main/scala/de/nowchess/bot/bots/nnue/NNUEWeights_v{version}.scala"
|
||||
|
||||
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
|
||||
|
||||
print("Available checkpoints:")
|
||||
for v in available:
|
||||
weights_file = f"nnue_weights_v{v}.pt"
|
||||
size = Path(weights_file).stat().st_size / (1024**2) # MB
|
||||
print(f" v{v} ({size:.1f} MB)")
|
||||
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())
|
||||
Reference in New Issue
Block a user