Initial commit

cleaning/refine.py

@@ -0,0 +1,259 @@
+#!/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()