Add stub
This commit is contained in:
@@ -1,168 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
"""
|
|
||||||
Turn refined PGN games (Site/WhiteElo/BlackElo headers + movetext with
|
|
||||||
%eval comments) into a CSV where each row is the board state *after* one
|
|
||||||
move.
|
|
||||||
|
|
||||||
Row layout (67 columns):
|
|
||||||
game_id, a8, b8, c8, d8, e8, f8, g8, h8, a7, ... , h1, eval, turn
|
|
||||||
|
|
||||||
- game_id : 0-based index of the game in the input file. Not asked for
|
|
||||||
explicitly, but with many games concatenated into one CSV
|
|
||||||
there is no other way to tell which rows belong to which
|
|
||||||
game, so I added it as column 1. Drop it with --no-game-id
|
|
||||||
if you don't want it.
|
|
||||||
- a8..h1 : the 64 squares (rank 8 down to rank 1, file a to h), each
|
|
||||||
holding the FEN piece letter on that square after the move
|
|
||||||
(uppercase = white, lowercase = black) or an empty string
|
|
||||||
if the square is empty.
|
|
||||||
- eval : the %eval attached to that move (blank if that particular
|
|
||||||
move had none).
|
|
||||||
- turn : "white" or "black" -- whose turn it is to move NEXT.
|
|
||||||
|
|
||||||
Streams game-by-game (like the two earlier scripts) so it stays cheap on
|
|
||||||
RAM no matter how large the input file is. Requires `pip install chess`.
|
|
||||||
|
|
||||||
Usage:
|
|
||||||
python pgn_to_csv.py refined.pgn out.csv
|
|
||||||
python pgn_to_csv.py refined.pgn out.csv --no-game-id
|
|
||||||
"""
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import csv
|
|
||||||
import re
|
|
||||||
import sys
|
|
||||||
|
|
||||||
import chess
|
|
||||||
|
|
||||||
TAG_RE = re.compile(r'^\[(\w+)\s+"(.*)"\]\s*$')
|
|
||||||
EVAL_RE = re.compile(r'\[%eval\s+([^\]]+)\]')
|
|
||||||
MOVE_NUM_RE = re.compile(r'^\d+\.+$')
|
|
||||||
RESULT_RE = re.compile(r'^(1-0|0-1|1/2-1/2|\*)$')
|
|
||||||
# Strip trailing annotation glyphs python-chess's SAN parser won't accept
|
|
||||||
# (e.g. "e5?!", "Nc3!!", "Qxd1+!") while keeping the "+"/"#" that ARE part
|
|
||||||
# of legal SAN.
|
|
||||||
SUFFIX_ANNOTATION_RE = re.compile(r'[!?]+$')
|
|
||||||
|
|
||||||
SQUARE_COLUMNS = [f + str(r) for r in range(8, 0, -1) for f in "abcdefgh"]
|
|
||||||
|
|
||||||
|
|
||||||
def tokenize_movetext(movetext: str):
|
|
||||||
"""Yield (san_move, eval_str_or_None) pairs from a movetext blob."""
|
|
||||||
tokens = re.findall(r'\{[^}]*\}|\S+', movetext)
|
|
||||||
i = 0
|
|
||||||
n = len(tokens)
|
|
||||||
while i < n:
|
|
||||||
tok = tokens[i]
|
|
||||||
if MOVE_NUM_RE.match(tok) or RESULT_RE.match(tok):
|
|
||||||
i += 1
|
|
||||||
continue
|
|
||||||
move_san = tok
|
|
||||||
eval_val = None
|
|
||||||
if i + 1 < n and tokens[i + 1].startswith('{'):
|
|
||||||
m = EVAL_RE.search(tokens[i + 1])
|
|
||||||
if m:
|
|
||||||
eval_val = m.group(1)
|
|
||||||
i += 2
|
|
||||||
else:
|
|
||||||
i += 1
|
|
||||||
yield move_san, eval_val
|
|
||||||
|
|
||||||
|
|
||||||
def board_row(board: chess.Board):
|
|
||||||
"""Return the 64 square values for the current board, a8..h1 order."""
|
|
||||||
row = []
|
|
||||||
for name in SQUARE_COLUMNS:
|
|
||||||
sq = chess.parse_square(name)
|
|
||||||
piece = board.piece_at(sq)
|
|
||||||
row.append(piece.symbol() if piece else "")
|
|
||||||
return row
|
|
||||||
|
|
||||||
|
|
||||||
def process_game(game_id, header_lines, movetext_lines, writer, include_game_id):
|
|
||||||
if not movetext_lines:
|
|
||||||
return 0
|
|
||||||
|
|
||||||
movetext = ' '.join(line.strip() for line in movetext_lines if line.strip())
|
|
||||||
|
|
||||||
board = chess.Board()
|
|
||||||
rows_written = 0
|
|
||||||
|
|
||||||
for move_san, eval_val in tokenize_movetext(movetext):
|
|
||||||
clean_san = SUFFIX_ANNOTATION_RE.sub('', move_san)
|
|
||||||
try:
|
|
||||||
board.push_san(clean_san)
|
|
||||||
except (chess.IllegalMoveError, chess.InvalidMoveError, chess.AmbiguousMoveError, ValueError) as e:
|
|
||||||
print(f" [game {game_id}] stopping early: could not parse move "
|
|
||||||
f"'{move_san}' ({e})", file=sys.stderr)
|
|
||||||
break
|
|
||||||
|
|
||||||
turn = "white" if board.turn == chess.WHITE else "black"
|
|
||||||
row = (board_row(board) + [eval_val if eval_val is not None else "", turn])
|
|
||||||
if include_game_id:
|
|
||||||
row = [game_id] + row
|
|
||||||
writer.writerow(row)
|
|
||||||
rows_written += 1
|
|
||||||
|
|
||||||
return rows_written
|
|
||||||
|
|
||||||
|
|
||||||
def refine(input_path: str, output_path: str, include_game_id: bool):
|
|
||||||
header_lines = []
|
|
||||||
movetext_lines = []
|
|
||||||
state = "seek_header"
|
|
||||||
game_id = 0
|
|
||||||
total_rows = 0
|
|
||||||
|
|
||||||
header_row = SQUARE_COLUMNS + ["eval", "turn"]
|
|
||||||
if include_game_id:
|
|
||||||
header_row = ["game_id"] + header_row
|
|
||||||
|
|
||||||
with open(input_path, "r", encoding="utf-8", errors="replace", buffering=1024 * 1024) as fin, \
|
|
||||||
open(output_path, "w", encoding="utf-8", newline="") as fout:
|
|
||||||
|
|
||||||
writer = csv.writer(fout)
|
|
||||||
writer.writerow(header_row)
|
|
||||||
|
|
||||||
for raw_line in fin:
|
|
||||||
line = raw_line.rstrip("\n")
|
|
||||||
|
|
||||||
if state == "seek_header":
|
|
||||||
if line.startswith("["):
|
|
||||||
header_lines.append(line)
|
|
||||||
elif line.strip() == "" and header_lines:
|
|
||||||
state = "seek_movetext"
|
|
||||||
elif state == "seek_movetext":
|
|
||||||
if line.strip() == "":
|
|
||||||
total_rows += process_game(game_id, header_lines, movetext_lines,
|
|
||||||
writer, include_game_id)
|
|
||||||
game_id += 1
|
|
||||||
header_lines = []
|
|
||||||
movetext_lines = []
|
|
||||||
state = "seek_header"
|
|
||||||
else:
|
|
||||||
movetext_lines.append(line)
|
|
||||||
|
|
||||||
if header_lines and movetext_lines:
|
|
||||||
total_rows += process_game(game_id, header_lines, movetext_lines,
|
|
||||||
writer, include_game_id)
|
|
||||||
game_id += 1
|
|
||||||
|
|
||||||
return game_id, total_rows
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
parser = argparse.ArgumentParser(
|
|
||||||
description="Convert refined PGN games into a per-move board-state CSV.")
|
|
||||||
parser.add_argument("input", help="Path to the refined PGN file")
|
|
||||||
parser.add_argument("output", help="Path to write the CSV file")
|
|
||||||
parser.add_argument("--no-game-id", action="store_true",
|
|
||||||
help="Omit the leading game_id column (66-col strict layout).")
|
|
||||||
args = parser.parse_args()
|
|
||||||
|
|
||||||
games, rows = refine(args.input, args.output, include_game_id=not args.no_game_id)
|
|
||||||
print(f"Processed {games} games -> {rows} rows written to {args.output}")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
5
requirements.txt
Normal file
5
requirements.txt
Normal file
@@ -0,0 +1,5 @@
|
|||||||
|
python-chess
|
||||||
|
|
||||||
|
torch
|
||||||
|
pandas
|
||||||
|
pandas-stubs
|
||||||
97
src/data_processing/dataset_loader.py
Normal file
97
src/data_processing/dataset_loader.py
Normal file
@@ -0,0 +1,97 @@
|
|||||||
|
import pandas as pd
|
||||||
|
import torch
|
||||||
|
|
||||||
|
INPUT_CSV = "../../dataset/output.csv"
|
||||||
|
OUTPUT_FILE = "../../dataset/output.pt"
|
||||||
|
|
||||||
|
SQUARES = [
|
||||||
|
f"{file}{rank}"
|
||||||
|
for rank in range(8, 0, -1)
|
||||||
|
for file in "abcdefgh"
|
||||||
|
]
|
||||||
|
|
||||||
|
PIECE_TO_PLANE = {
|
||||||
|
"P": 0,
|
||||||
|
"N": 1,
|
||||||
|
"B": 2,
|
||||||
|
"R": 3,
|
||||||
|
"Q": 4,
|
||||||
|
"K": 5,
|
||||||
|
"p": 6,
|
||||||
|
"n": 7,
|
||||||
|
"b": 8,
|
||||||
|
"r": 9,
|
||||||
|
"q": 10,
|
||||||
|
"k": 11,
|
||||||
|
}
|
||||||
|
|
||||||
|
def board_to_planes(row):
|
||||||
|
planes = torch.zeros((12, 8, 8), dtype=torch.float32)
|
||||||
|
|
||||||
|
for index, square in enumerate(SQUARES):
|
||||||
|
piece = row[square]
|
||||||
|
|
||||||
|
if pd.isna(piece):
|
||||||
|
continue
|
||||||
|
|
||||||
|
plane = PIECE_TO_PLANE.get(piece)
|
||||||
|
|
||||||
|
if plane is None:
|
||||||
|
continue
|
||||||
|
|
||||||
|
row_index = index // 8
|
||||||
|
col_index = index % 8
|
||||||
|
|
||||||
|
planes[plane, row_index, col_index] = 1.0
|
||||||
|
|
||||||
|
return planes
|
||||||
|
|
||||||
|
def parse_eval(value):
|
||||||
|
value = str(value).strip()
|
||||||
|
|
||||||
|
if value.startswith("#"):
|
||||||
|
mate_distance = int(value[1:])
|
||||||
|
|
||||||
|
if mate_distance > 0:
|
||||||
|
return 10.0
|
||||||
|
else:
|
||||||
|
return -10.0
|
||||||
|
|
||||||
|
return float(value)
|
||||||
|
|
||||||
|
def main():
|
||||||
|
df = pd.read_csv(INPUT_CSV)
|
||||||
|
|
||||||
|
positions = []
|
||||||
|
evaluations = []
|
||||||
|
turns = []
|
||||||
|
|
||||||
|
for _, row in df.iterrows():
|
||||||
|
positions.append(board_to_planes(row))
|
||||||
|
|
||||||
|
evaluations.append(parse_eval(row["eval"]))
|
||||||
|
|
||||||
|
turns.append(1 if row["turn"] == "white" else 0)
|
||||||
|
|
||||||
|
|
||||||
|
X = torch.stack(positions)
|
||||||
|
y = torch.tensor(evaluations, dtype=torch.float32)
|
||||||
|
turn = torch.tensor(turns, dtype=torch.float32)
|
||||||
|
|
||||||
|
torch.save(
|
||||||
|
{
|
||||||
|
"X": X,
|
||||||
|
"y": y,
|
||||||
|
"turn": turn,
|
||||||
|
},
|
||||||
|
OUTPUT_FILE,
|
||||||
|
)
|
||||||
|
|
||||||
|
print(f"Loaded {len(df):,} positions")
|
||||||
|
print(f"X shape: {X.shape}")
|
||||||
|
print(f"y shape: {y.shape}")
|
||||||
|
print(f"turn shape: {turn.shape}")
|
||||||
|
print(f"Saved to: {OUTPUT_FILE}")
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
345
src/data_processing/format.py
Normal file
345
src/data_processing/format.py
Normal file
@@ -0,0 +1,345 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""
|
||||||
|
Turn refined PGN games into a CSV where each row is the board state
|
||||||
|
after one evaluated move.
|
||||||
|
|
||||||
|
Row layout:
|
||||||
|
game_id, a8, b8, ..., h1, eval, turn
|
||||||
|
|
||||||
|
- game_id : Lichess game ID extracted from the Site header.
|
||||||
|
Example:
|
||||||
|
[Site "https://lichess.org/qE0J8lxe"]
|
||||||
|
becomes:
|
||||||
|
qE0J8lxe
|
||||||
|
|
||||||
|
- a8..h1 : the 64 squares after the move.
|
||||||
|
Uppercase = white, lowercase = black,
|
||||||
|
empty string = empty square.
|
||||||
|
|
||||||
|
- eval : the %eval attached to the move.
|
||||||
|
|
||||||
|
- turn : "white" or "black", meaning whose turn it is NEXT.
|
||||||
|
|
||||||
|
Moves without an evaluation are skipped.
|
||||||
|
|
||||||
|
Streams game-by-game, so memory usage stays low.
|
||||||
|
|
||||||
|
Requires:
|
||||||
|
pip install chess
|
||||||
|
|
||||||
|
Usage:
|
||||||
|
python pgn_to_csv.py refined.pgn out.csv
|
||||||
|
python pgn_to_csv.py refined.pgn out.csv --no-game-id
|
||||||
|
"""
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import csv
|
||||||
|
import re
|
||||||
|
import sys
|
||||||
|
|
||||||
|
import chess
|
||||||
|
|
||||||
|
TAG_RE = re.compile(r'^\[(\w+)\s+"(.*)"\]\s*$')
|
||||||
|
EVAL_RE = re.compile(r'\[%eval\s+([^\]]+)\]')
|
||||||
|
MOVE_NUM_RE = re.compile(r'^\d+\.+$')
|
||||||
|
RESULT_RE = re.compile(r'^(1-0|0-1|1/2-1/2|\*)$')
|
||||||
|
SUFFIX_ANNOTATION_RE = re.compile(r'[!?]+$')
|
||||||
|
|
||||||
|
SQUARE_COLUMNS = [
|
||||||
|
f + str(r)
|
||||||
|
for r in range(8, 0, -1)
|
||||||
|
for f in "abcdefgh"
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def parse_headers(header_lines):
|
||||||
|
"""
|
||||||
|
Convert PGN header lines into a dictionary.
|
||||||
|
|
||||||
|
Example:
|
||||||
|
[Site "https://lichess.org/qE0J8lxe"]
|
||||||
|
|
||||||
|
becomes:
|
||||||
|
{
|
||||||
|
"Site": "https://lichess.org/qE0J8lxe"
|
||||||
|
}
|
||||||
|
"""
|
||||||
|
headers = {}
|
||||||
|
|
||||||
|
for line in header_lines:
|
||||||
|
match = TAG_RE.match(line)
|
||||||
|
|
||||||
|
if match:
|
||||||
|
key, value = match.groups()
|
||||||
|
headers[key] = value
|
||||||
|
|
||||||
|
return headers
|
||||||
|
|
||||||
|
|
||||||
|
def extract_game_id(site):
|
||||||
|
"""
|
||||||
|
Extract the Lichess game ID from the Site URL.
|
||||||
|
|
||||||
|
Example:
|
||||||
|
https://lichess.org/qE0J8lxe
|
||||||
|
-> qE0J8lxe
|
||||||
|
"""
|
||||||
|
if not site:
|
||||||
|
return None
|
||||||
|
|
||||||
|
return site.rstrip("/").split("/")[-1]
|
||||||
|
|
||||||
|
|
||||||
|
def tokenize_movetext(movetext):
|
||||||
|
"""
|
||||||
|
Yield (san_move, eval_str_or_None) pairs from a movetext blob.
|
||||||
|
"""
|
||||||
|
tokens = re.findall(r'\{[^}]*\}|\S+', movetext)
|
||||||
|
|
||||||
|
i = 0
|
||||||
|
n = len(tokens)
|
||||||
|
|
||||||
|
while i < n:
|
||||||
|
tok = tokens[i]
|
||||||
|
|
||||||
|
if MOVE_NUM_RE.match(tok) or RESULT_RE.match(tok):
|
||||||
|
i += 1
|
||||||
|
continue
|
||||||
|
|
||||||
|
move_san = tok
|
||||||
|
eval_val = None
|
||||||
|
|
||||||
|
if i + 1 < n and tokens[i + 1].startswith('{'):
|
||||||
|
match = EVAL_RE.search(tokens[i + 1])
|
||||||
|
|
||||||
|
if match:
|
||||||
|
eval_val = match.group(1)
|
||||||
|
|
||||||
|
i += 2
|
||||||
|
|
||||||
|
else:
|
||||||
|
i += 1
|
||||||
|
|
||||||
|
yield move_san, eval_val
|
||||||
|
|
||||||
|
|
||||||
|
def board_row(board):
|
||||||
|
"""
|
||||||
|
Return the 64 square values for the current board,
|
||||||
|
in a8..h1 order.
|
||||||
|
"""
|
||||||
|
row = []
|
||||||
|
|
||||||
|
for name in SQUARE_COLUMNS:
|
||||||
|
square = chess.parse_square(name)
|
||||||
|
piece = board.piece_at(square)
|
||||||
|
|
||||||
|
row.append(piece.symbol() if piece else "")
|
||||||
|
|
||||||
|
return row
|
||||||
|
|
||||||
|
|
||||||
|
def process_game(header_lines, movetext_lines, writer, include_game_id):
|
||||||
|
"""
|
||||||
|
Process one game and write its evaluated positions.
|
||||||
|
|
||||||
|
Returns:
|
||||||
|
number of rows written
|
||||||
|
"""
|
||||||
|
|
||||||
|
if not movetext_lines:
|
||||||
|
return 0
|
||||||
|
|
||||||
|
headers = parse_headers(header_lines)
|
||||||
|
|
||||||
|
site = headers.get("Site")
|
||||||
|
game_id = extract_game_id(site)
|
||||||
|
|
||||||
|
if game_id is None:
|
||||||
|
print(
|
||||||
|
" [warning] Game has no valid Site header, skipping.",
|
||||||
|
file=sys.stderr
|
||||||
|
)
|
||||||
|
return 0
|
||||||
|
|
||||||
|
movetext = " ".join(
|
||||||
|
line.strip()
|
||||||
|
for line in movetext_lines
|
||||||
|
if line.strip()
|
||||||
|
)
|
||||||
|
|
||||||
|
board = chess.Board()
|
||||||
|
rows_written = 0
|
||||||
|
|
||||||
|
for move_san, eval_val in tokenize_movetext(movetext):
|
||||||
|
|
||||||
|
# No evaluation means this position has no target value.
|
||||||
|
# Do not put it into the regression dataset.
|
||||||
|
if eval_val is None:
|
||||||
|
continue
|
||||||
|
|
||||||
|
clean_san = SUFFIX_ANNOTATION_RE.sub("", move_san)
|
||||||
|
|
||||||
|
try:
|
||||||
|
board.push_san(clean_san)
|
||||||
|
|
||||||
|
except (
|
||||||
|
chess.IllegalMoveError,
|
||||||
|
chess.InvalidMoveError,
|
||||||
|
chess.AmbiguousMoveError,
|
||||||
|
ValueError,
|
||||||
|
) as e:
|
||||||
|
|
||||||
|
print(
|
||||||
|
f" [game {game_id}] stopping early: "
|
||||||
|
f"could not parse move '{move_san}' ({e})",
|
||||||
|
file=sys.stderr,
|
||||||
|
)
|
||||||
|
|
||||||
|
break
|
||||||
|
|
||||||
|
turn = (
|
||||||
|
"white"
|
||||||
|
if board.turn == chess.WHITE
|
||||||
|
else "black"
|
||||||
|
)
|
||||||
|
|
||||||
|
row = board_row(board)
|
||||||
|
|
||||||
|
row += [
|
||||||
|
eval_val,
|
||||||
|
turn,
|
||||||
|
]
|
||||||
|
|
||||||
|
if include_game_id:
|
||||||
|
row = [game_id] + row
|
||||||
|
|
||||||
|
writer.writerow(row)
|
||||||
|
|
||||||
|
rows_written += 1
|
||||||
|
|
||||||
|
return rows_written
|
||||||
|
|
||||||
|
|
||||||
|
def refine(input_path, output_path, include_game_id):
|
||||||
|
header_lines = []
|
||||||
|
movetext_lines = []
|
||||||
|
|
||||||
|
state = "seek_header"
|
||||||
|
|
||||||
|
games = 0
|
||||||
|
total_rows = 0
|
||||||
|
|
||||||
|
header_row = SQUARE_COLUMNS + ["eval", "turn"]
|
||||||
|
|
||||||
|
if include_game_id:
|
||||||
|
header_row = ["game_id"] + header_row
|
||||||
|
|
||||||
|
with (
|
||||||
|
open(
|
||||||
|
input_path,
|
||||||
|
"r",
|
||||||
|
encoding="utf-8",
|
||||||
|
errors="replace",
|
||||||
|
buffering=1024 * 1024,
|
||||||
|
) as fin,
|
||||||
|
|
||||||
|
open(
|
||||||
|
output_path,
|
||||||
|
"w",
|
||||||
|
encoding="utf-8",
|
||||||
|
newline="",
|
||||||
|
) as fout
|
||||||
|
):
|
||||||
|
|
||||||
|
writer = csv.writer(fout)
|
||||||
|
writer.writerow(header_row)
|
||||||
|
|
||||||
|
for raw_line in fin:
|
||||||
|
line = raw_line.rstrip("\n")
|
||||||
|
|
||||||
|
if state == "seek_header":
|
||||||
|
|
||||||
|
if line.startswith("["):
|
||||||
|
header_lines.append(line)
|
||||||
|
|
||||||
|
elif line.strip() == "" and header_lines:
|
||||||
|
state = "seek_movetext"
|
||||||
|
|
||||||
|
elif state == "seek_movetext":
|
||||||
|
|
||||||
|
if line.strip() == "":
|
||||||
|
rows = process_game(
|
||||||
|
header_lines,
|
||||||
|
movetext_lines,
|
||||||
|
writer,
|
||||||
|
include_game_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
total_rows += rows
|
||||||
|
games += 1
|
||||||
|
|
||||||
|
header_lines = []
|
||||||
|
movetext_lines = []
|
||||||
|
|
||||||
|
state = "seek_header"
|
||||||
|
|
||||||
|
else:
|
||||||
|
movetext_lines.append(line)
|
||||||
|
|
||||||
|
# Handle final game if the file doesn't end with a blank line.
|
||||||
|
if header_lines and movetext_lines:
|
||||||
|
rows = process_game(
|
||||||
|
header_lines,
|
||||||
|
movetext_lines,
|
||||||
|
writer,
|
||||||
|
include_game_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
total_rows += rows
|
||||||
|
games += 1
|
||||||
|
|
||||||
|
return games, total_rows
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
parser = argparse.ArgumentParser(
|
||||||
|
description=(
|
||||||
|
"Convert refined PGN games into a "
|
||||||
|
"per-move board-state CSV."
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument(
|
||||||
|
"input",
|
||||||
|
help="Path to the refined PGN file",
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument(
|
||||||
|
"output",
|
||||||
|
help="Path to write the CSV file",
|
||||||
|
)
|
||||||
|
|
||||||
|
parser.add_argument(
|
||||||
|
"--no-game-id",
|
||||||
|
action="store_true",
|
||||||
|
help="Omit the leading game_id column.",
|
||||||
|
)
|
||||||
|
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
games, rows = refine(
|
||||||
|
args.input,
|
||||||
|
args.output,
|
||||||
|
include_game_id=not args.no_game_id,
|
||||||
|
)
|
||||||
|
|
||||||
|
print(
|
||||||
|
f"Processed {games:,} games "
|
||||||
|
f"-> {rows:,} evaluated positions written to "
|
||||||
|
f"{args.output}"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
main()
|
||||||
71
src/training/decision_tree/main.py
Normal file
71
src/training/decision_tree/main.py
Normal file
@@ -0,0 +1,71 @@
|
|||||||
|
import os
|
||||||
|
|
||||||
|
from sklearn.ensemble import HistGradientBoostingRegressor, GradientBoostingRegressor, RandomForestRegressor
|
||||||
|
from sklearn.model_selection import train_test_split, GroupShuffleSplit
|
||||||
|
from sklearn.tree import DecisionTreeRegressor
|
||||||
|
|
||||||
|
from src.training.decision_tree import preprocessor
|
||||||
|
from src.training.decision_tree.model import Model
|
||||||
|
from src.training.decision_tree.visualizer import Visualizer
|
||||||
|
|
||||||
|
|
||||||
|
def create_split(x, y, game_ids):
|
||||||
|
|
||||||
|
splitter = GroupShuffleSplit(
|
||||||
|
n_splits=1,
|
||||||
|
test_size=0.10,
|
||||||
|
random_state=42
|
||||||
|
)
|
||||||
|
|
||||||
|
train_indices, test_indices = next(splitter.split(x, y, groups=game_ids))
|
||||||
|
|
||||||
|
x_train = x[train_indices]
|
||||||
|
x_test = x[test_indices]
|
||||||
|
|
||||||
|
y_train = y[train_indices]
|
||||||
|
y_test = y[test_indices]
|
||||||
|
|
||||||
|
train_games = set(game_ids[train_indices])
|
||||||
|
test_games = set(game_ids[test_indices])
|
||||||
|
|
||||||
|
print("Train positions:", len(x_train))
|
||||||
|
print("Test positions:", len(x_test))
|
||||||
|
print("Train games:", len(train_games))
|
||||||
|
print("Test games:", len(test_games))
|
||||||
|
print("Overlapping games:", len(train_games & test_games))
|
||||||
|
|
||||||
|
return x_train, x_test, y_train, y_test
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
|
||||||
|
model_path = "C:/Users/Laurent/Desktop/chess-engine/src/training/decision_tree/cache/model.joblib"
|
||||||
|
|
||||||
|
if os.path.exists(model_path):
|
||||||
|
model = Model.load(model_path)
|
||||||
|
|
||||||
|
else:
|
||||||
|
x, y, game_ids = preprocessor.run("C:/Users/Laurent/Desktop/chess-engine/dataset/v1/output.csv")
|
||||||
|
|
||||||
|
split = create_split(x, y, game_ids)
|
||||||
|
|
||||||
|
depth = 35
|
||||||
|
leafs = int(depth * 3.5)
|
||||||
|
|
||||||
|
model = Model(
|
||||||
|
tree=RandomForestRegressor(
|
||||||
|
n_estimators=25,
|
||||||
|
max_depth=depth,
|
||||||
|
max_features=0.35,
|
||||||
|
min_samples_leaf=leafs,
|
||||||
|
random_state=42
|
||||||
|
),
|
||||||
|
feature_names=preprocessor.get_feature_names(),
|
||||||
|
train_test_split=split
|
||||||
|
)
|
||||||
|
|
||||||
|
model.train()
|
||||||
|
model.evaluate()
|
||||||
|
#model.save_model(model_path)
|
||||||
|
|
||||||
|
visualizer = Visualizer(model=model)
|
||||||
|
visualizer.show_feature_importance(15)
|
||||||
75
src/training/decision_tree/model.py
Normal file
75
src/training/decision_tree/model.py
Normal file
@@ -0,0 +1,75 @@
|
|||||||
|
import os
|
||||||
|
import joblib
|
||||||
|
from sklearn.metrics import mean_absolute_error, mean_squared_error
|
||||||
|
|
||||||
|
|
||||||
|
class Model:
|
||||||
|
def __init__(self, tree=None, feature_names=None, train_test_split=None):
|
||||||
|
self._model = tree
|
||||||
|
self._feature_names = feature_names
|
||||||
|
self._x_train = self._x_test = None
|
||||||
|
self._y_train = self._y_test = None
|
||||||
|
|
||||||
|
if train_test_split is not None:
|
||||||
|
self._x_train, self._x_test, \
|
||||||
|
self._y_train, self._y_test = train_test_split
|
||||||
|
|
||||||
|
@classmethod
|
||||||
|
def load(cls, path):
|
||||||
|
data = joblib.load(path)
|
||||||
|
|
||||||
|
model = cls()
|
||||||
|
model._model = data["model"]
|
||||||
|
model._feature_names = data["feature_names"]
|
||||||
|
|
||||||
|
print(f"Loaded model from {path}")
|
||||||
|
|
||||||
|
return model
|
||||||
|
|
||||||
|
def get_feature_names(self):
|
||||||
|
return self._feature_names
|
||||||
|
|
||||||
|
def get_feature_importance(self):
|
||||||
|
importances = self._model.feature_importances_
|
||||||
|
|
||||||
|
return sorted(
|
||||||
|
zip(self._feature_names, importances),
|
||||||
|
key=lambda x: x[1],
|
||||||
|
reverse=True
|
||||||
|
)
|
||||||
|
|
||||||
|
def train(self):
|
||||||
|
print("Start training...")
|
||||||
|
self._model.fit(self._x_train, self._y_train)
|
||||||
|
print("Training complete.")
|
||||||
|
|
||||||
|
def evaluate(self):
|
||||||
|
print("Evaluating model...")
|
||||||
|
|
||||||
|
train_predictions = self._model.predict(self._x_train)
|
||||||
|
test_predictions = self._model.predict(self._x_test)
|
||||||
|
|
||||||
|
train_mae = mean_absolute_error(self._y_train, train_predictions)
|
||||||
|
train_mse = mean_squared_error(self._y_train, train_predictions)
|
||||||
|
|
||||||
|
test_mae = mean_absolute_error(self._y_test, test_predictions)
|
||||||
|
test_mse = mean_squared_error(self._y_test, test_predictions)
|
||||||
|
|
||||||
|
print()
|
||||||
|
print("Training:")
|
||||||
|
print("MAE:", train_mae)
|
||||||
|
print("MSE:", train_mse)
|
||||||
|
|
||||||
|
print()
|
||||||
|
print("Test:")
|
||||||
|
print("MAE:", test_mae)
|
||||||
|
print("MSE:", test_mse)
|
||||||
|
|
||||||
|
|
||||||
|
def save_model(self, path):
|
||||||
|
joblib.dump({
|
||||||
|
"model": self._model,
|
||||||
|
"feature_names": self._feature_names
|
||||||
|
}, path)
|
||||||
|
|
||||||
|
print(f"Saved model to {path}")
|
||||||
1030
src/training/decision_tree/preprocessor.py
Normal file
1030
src/training/decision_tree/preprocessor.py
Normal file
File diff suppressed because it is too large
Load Diff
50
src/training/decision_tree/visualizer.py
Normal file
50
src/training/decision_tree/visualizer.py
Normal file
@@ -0,0 +1,50 @@
|
|||||||
|
import chess
|
||||||
|
import numpy as np
|
||||||
|
from matplotlib import pyplot as plt
|
||||||
|
|
||||||
|
|
||||||
|
class Visualizer:
|
||||||
|
def __init__(self, model):
|
||||||
|
self._model = model
|
||||||
|
|
||||||
|
def show_feature_importance(self, top_n=15):
|
||||||
|
importances = self._model.get_feature_importance()
|
||||||
|
print("Feature importance")
|
||||||
|
print("-------------------------")
|
||||||
|
|
||||||
|
for feature_name, importance in importances[:top_n]:
|
||||||
|
print(f"{feature_name:<35} {importance:.6f}")
|
||||||
|
|
||||||
|
def show_piece_heatmaps(self):
|
||||||
|
importances = dict(self._model.get_feature_importance())
|
||||||
|
|
||||||
|
pieces = [
|
||||||
|
"P", "N", "B", "R", "Q", "K",
|
||||||
|
"p", "n", "b", "r", "q", "k"
|
||||||
|
]
|
||||||
|
|
||||||
|
fig, axes = plt.subplots(3, 4, figsize=(16, 12))
|
||||||
|
|
||||||
|
for ax, piece in zip(axes.flat, pieces):
|
||||||
|
board = np.zeros((8, 8))
|
||||||
|
|
||||||
|
for row, rank in enumerate(range(7, -1, -1)):
|
||||||
|
for col, file in enumerate(range(8)):
|
||||||
|
square = chess.square_name(chess.square(file, rank))
|
||||||
|
feature_name = f"{piece}_{square}"
|
||||||
|
board[row, col] = importances.get(feature_name, 0)
|
||||||
|
|
||||||
|
image = ax.imshow(board)
|
||||||
|
|
||||||
|
ax.set_xticks(range(8))
|
||||||
|
ax.set_xticklabels("abcdefgh")
|
||||||
|
ax.set_yticks(range(8))
|
||||||
|
ax.set_yticklabels("87654321")
|
||||||
|
ax.set_title(piece)
|
||||||
|
|
||||||
|
fig.colorbar(image, ax=ax, fraction=0.046, pad=0.04)
|
||||||
|
|
||||||
|
fig.suptitle("Piece-square feature importance", fontsize=16)
|
||||||
|
plt.tight_layout()
|
||||||
|
plt.show()
|
||||||
|
input()
|
||||||
Reference in New Issue
Block a user