rework dataset function and class structure

server/dataset/DatasetCurator.py

@@ -1,6 +1,8 @@
-import argparse, hashlib, shutil, glob, json
+import argparse, hashlib, shutil, json
 from pathlib import Path
 
+from server.dataset.DatasetIO import DatasetIO
+
 class DatasetCurator:
   def __init__(self, input_files:list[str], output_file:str, min_turn:int=6, late_turn:int=20, max_safe_options:int=2, min_score:int=3, append:bool=False, archive_input:bool=False, archive_dir:str|None=None):
     self.input_files = input_files
@@ -11,45 +13,7 @@ class DatasetCurator:
     self.min_score = min_score
     self.append = append
     self.archive_input = archive_input
-    self.archive_dir = (
-      Path(archive_dir) if archive_dir else self.output_file.parent / "archive"
-    )
-
-  def _resolve_input_files(self):
-    resolved = []
-    seen = set()
-
-    for item in self.input_files:
-      path = Path(item)
-      if path.is_dir():
-        for file_path in sorted(path.rglob("*.jsonl")):
-          key = str(file_path.resolve())
-          if key in seen:
-            continue
-          seen.add(key)
-          resolved.append(file_path)
-        continue
-
-      if any(ch in item for ch in "*?[]"):
-        for match in sorted(glob.glob(item)):
-          file_path = Path(match)
-          if not file_path.is_file():
-            continue
-          key = str(file_path.resolve())
-          if key in seen:
-            continue
-          seen.add(key)
-          resolved.append(file_path)
-        continue
-
-      if path.is_file():
-        key = str(path.resolve())
-        if key in seen:
-          continue
-        seen.add(key)
-        resolved.append(path)
-
-    return resolved
+    self.archive_dir = Path(archive_dir) if archive_dir else self.output_file.parent / "archive"
 
   def _safe_options_count(self, row:dict):
     history = row.get("history", {})
@@ -98,7 +62,7 @@ class DatasetCurator:
 
   def curate(self):
     self.output_file.parent.mkdir(parents=True, exist_ok=True)
-    input_paths = self._resolve_input_files()
+    input_paths = DatasetIO.resolve_input_files(self.input_files)
 
     total = 0
     kept = 0
@@ -108,56 +72,47 @@ class DatasetCurator:
     seen_states = set()
 
     if self.append and self.output_file.exists():
-      with self.output_file.open("r", encoding="utf-8") as existing:
-        for line in existing:
-          if not line.strip():
-            continue
-          row = json.loads(line)
-          state_key = self._state_hash(row)
-          seen_states.add((state_key, row.get("move")))
+      for row in DatasetIO.iter_jsonl_rows(self.output_file):
+        state_key = self._state_hash(row)
+        seen_states.add((state_key, row.get("move")))
 
     mode = "a" if self.append else "w"
-    with self.output_file.open(mode, encoding="utf-8") as dst:
+    with DatasetIO.open_text(self.output_file, mode) as dst:
       for input_path in input_paths:
-        with input_path.open("r", encoding="utf-8") as src:
-          for line in src:
-            if not line.strip():
-              continue
+        for row in DatasetIO.iter_jsonl_rows(input_path):
+          total += 1
 
-            total += 1
-            row = json.loads(line)
+          if not row.get("is_good_move", False):
+            skipped_quality += 1
+            continue
 
-            if not row.get("is_good_move", False):
-              skipped_quality += 1
-              continue
+          if int(row.get("turn", 0)) < self.min_turn:
+            skipped_turn += 1
+            continue
 
-            if int(row.get("turn", 0)) < self.min_turn:
-              skipped_turn += 1
-              continue
+          quality_score, safe_options = self._score(row)
+          if quality_score < self.min_score:
+            skipped_quality += 1
+            continue
 
-            quality_score, safe_options = self._score(row)
-            if quality_score < self.min_score:
-              skipped_quality += 1
-              continue
+          state_key = self._state_hash(row)
+          dedupe_key = (state_key, row.get("move"))
+          if dedupe_key in seen_states:
+            skipped_duplicate += 1
+            continue
+          seen_states.add(dedupe_key)
 
-            state_key = self._state_hash(row)
-            dedupe_key = (state_key, row.get("move"))
-            if dedupe_key in seen_states:
-              skipped_duplicate += 1
-              continue
-            seen_states.add(dedupe_key)
-
-            compact_row = {
-              "game_id": row.get("game_id"),
-              "turn": row.get("turn"),
-              "move": row.get("move"),
-              "game_type": row.get("game_type"),
-              "quality_score": quality_score,
-              "safe_options": safe_options,
-              "game_board": row.get("game_board"),
-            }
-            dst.write(json.dumps(compact_row, ensure_ascii=False) + "\n")
-            kept += 1
+          compact_row = {
+            "game_id": row.get("game_id"),
+            "turn": row.get("turn"),
+            "move": row.get("move"),
+            "game_type": row.get("game_type"),
+            "quality_score": quality_score,
+            "safe_options": safe_options,
+            "game_board": row.get("game_board"),
+          }
+          dst.write(json.dumps(compact_row, ensure_ascii=False) + "\n")
+          kept += 1
 
     archived_files = []
     if self.archive_input:
@@ -220,7 +175,7 @@ if __name__ == "__main__":
     "--input",
     action="append",
     required=True,
-    help="Input JSONL file, directory, or glob pattern. Repeat for multiple inputs.",
+    help="Input JSONL/JSONL.GZ file, directory, or glob pattern. Repeat for multiple inputs.",
   )
   parser.add_argument("--output", required=True, help="Output JSONL file")
   parser.add_argument("--min-turn", type=int, default=6)