chess-engine/src/data_processing/refine.py

#!/usr/bin/env python3
"""
Refine a huge Lichess-style PGN file, streaming one game at a time so RAM
stays flat regardless of file size (works fine on a 200GB file).

Rules applied to every game:
  1. If the game has no %eval annotations at all, drop it entirely.
  2. Keep only these header tags: Site, WhiteElo, BlackElo.
  3. In the movetext, remove every comment/annotation except %eval.
  4. Keep only the %eval annotation from comments containing multiple
     attributes, e.g. { [%eval 0.18] [%clk 0:01:00] } -> { [%eval 0.18] }.
  5. Remove non-%eval textual comments, e.g. { Black offers draw }.
  6. Remove NAGs such as !, ?, !?, ?!, etc.

Games are separated by a blank line after the movetext.

Usage:
    python refine_pgn.py input.pgn output.pgn
"""

import argparse
import re

# Only these header tags survive
KEEP_TAGS = {"Site", "WhiteElo", "BlackElo"}

# Matches a header tag line, e.g. [WhiteElo "1500"]
TAG_RE = re.compile(r'^\[(\w+)\s+"(.*)"\]\s*$')

# Matches a %eval annotation inside a PGN comment.
#
# Examples:
#   [%eval 0.18]
#   [%eval -1.25]
#   [%eval #26]
#
# The value is captured without the surrounding brackets.
EVAL_RE = re.compile(r'\[%eval\s+([^\]]+)\]')

# Matches PGN comments enclosed in { ... }.
COMMENT_RE = re.compile(r'\{([^}]*)\}')

# Detect presence of an %eval annotation anywhere in the movetext.
EVAL_PRESENT_RE = re.compile(r'\[%eval\b')


def clean_comment(match: re.Match) -> str:
    """
    Replace a PGN comment with only its %eval annotation.

    Examples:

        { [%eval 0.18] }
            -> { [%eval 0.18] }

        { [%eval 0.18] [%clk 0:01:00] }
            -> { [%eval 0.18] }

        { Black offers draw }
            -> ""

        { [%clk 0:01:00] }
            -> ""
    """
    comment = match.group(1)

    eval_match = EVAL_RE.search(comment)

    if eval_match:
        return "{ [%eval " + eval_match.group(1).strip() + "] }"

    return ""


def clean_movetext(movetext: str) -> str:
    """
    Remove every PGN comment except %eval annotations.

    NAGs such as !, ?, !?, ?! are also removed.
    """

    # Keep only %eval from {...} comments.
    text = COMMENT_RE.sub(clean_comment, movetext)

    # Remove NAGs / move annotations:
    #
    #   !
    #   ?
    #   !!
    #   ??
    #   !?
    #   ?!
    #
    # These are annotations, not %eval attributes.
    text = re.sub(r'(?<=[a-zA-Z0-9+#=])(?:!!|\?\?|!\?|\?!|!|\?)(?=\s|$)', '', text)

    # Collapse whitespace.
    text = re.sub(r'[ \t]+', ' ', text)

    # Remove excessive whitespace around eval comments.
    text = re.sub(r'\s+\{', ' {', text)
    text = re.sub(r'\}\s+', '} ', text)

    return text.strip()


def process_game(header_lines, movetext_lines, fout) -> bool:
    """
    Filter/transform one game and write it out if it qualifies.

    Returns True if the game was written.
    """

    if not movetext_lines:
        return False

    movetext = ' '.join(
        line.strip()
        for line in movetext_lines
        if line.strip()
    )

    # Drop games with no %eval at all.
    if not EVAL_PRESENT_RE.search(movetext):
        return False

    # Keep only whitelisted header tags.
    kept_headers = []

    for line in header_lines:
        m = TAG_RE.match(line)

        if m and m.group(1) in KEEP_TAGS:
            kept_headers.append(line.rstrip('\n'))

    # Remove everything except %eval.
    cleaned_movetext = clean_movetext(movetext)

    for h in kept_headers:
        fout.write(h + "\n")

    fout.write("\n")
    fout.write(cleaned_movetext + "\n")
    fout.write("\n")

    return True


def refine(input_path: str, output_path: str):
    games_in = 0
    games_out = 0

    header_lines = []
    movetext_lines = []

    # State:
    #   "seek_header"    -> collecting [Tag "..."] lines
    #   "seek_movetext"  -> collecting movetext until blank line
    state = "seek_header"

    with open(
        input_path,
        "r",
        encoding="utf-8",
        errors="replace",
        buffering=1024 * 1024
    ) as fin, open(
        output_path,
        "w",
        encoding="utf-8",
        buffering=1024 * 1024
    ) as fout:

        for raw_line in fin:
            line = raw_line.rstrip("\n")

            if state == "seek_header":

                if line.startswith("["):
                    header_lines.append(line)

                elif line.strip() == "":
                    # Blank line between headers and movetext.
                    if header_lines:
                        state = "seek_movetext"

                # Anything else while no headers are collected:
                # ignore it.

            elif state == "seek_movetext":

                if line.strip() == "":
                    # End of this game's movetext.
                    games_in += 1

                    if process_game(
                        header_lines,
                        movetext_lines,
                        fout
                    ):
                        games_out += 1

                    header_lines = []
                    movetext_lines = []
                    state = "seek_header"

                else:
                    movetext_lines.append(line)

        # Handle a file without a trailing blank line.
        if header_lines and movetext_lines:
            games_in += 1

            if process_game(
                header_lines,
                movetext_lines,
                fout
            ):
                games_out += 1

    return games_in, games_out


def main():
    parser = argparse.ArgumentParser(
        description=(
            "Refine a large PGN file: keep only selected headers "
            "and %eval annotations."
        )
    )

    parser.add_argument(
        "input",
        help="Path to the input PGN file"
    )

    parser.add_argument(
        "output",
        help="Path to write the refined PGN file"
    )

    args = parser.parse_args()

    total_in, total_out = refine(
        args.input,
        args.output
    )

    print(
        f"Read {total_in} games, "
        f"kept {total_out} (had %eval), "
        f"dropped {total_in - total_out} (no %eval)."
    )

    print(f"Output written to {args.output}")


if __name__ == "__main__":
    main()