Upload stratego/game_logger.py with huggingface_hub

stratego/game_logger.py

@@ -0,0 +1,134 @@
+# stratego/game_logger.py
+from __future__ import annotations
+import csv, os, datetime
+from typing import Optional, List
+
+
+class GameLogger:
+    """
+    One CSV file per game stored in logs/games/ folder.
+    
+    CSV Fields for training:
+      - turn, player, model_name
+      - move, from_pos, to_pos, piece_type
+      - board_state, available_moves, move_direction
+      - target_piece, battle_outcome
+      - prompt_name
+      - game_winner, game_result (filled post-game)
+    """
+    def __init__(self, out_dir: str, game_id: Optional[str] = None, prompt_name: str = "", game_type: str = "standard", board_size: int = 10):
+        # Create games subfolder
+        games_dir = os.path.join(out_dir, "games")
+        os.makedirs(games_dir, exist_ok=True)
+        
+        ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        self.game_id = game_id or ts
+        self.prompt_name = prompt_name
+        self.game_type = game_type  # "standard", "duel", or "custom"
+        self.board_size = board_size
+        self.path = os.path.join(games_dir, f"{self.game_id}.csv")
+        self._f = open(self.path, "w", newline="", encoding="utf-8")
+        # Use full set of fields up-front so partial logs are readable during play
+        self._writer = csv.DictWriter(
+            self._f,
+            fieldnames=[
+                "turn", "player", "model_name",
+                "move", "from_pos", "to_pos", "piece_type",
+                "board_state", "available_moves", "move_direction",
+                "target_piece", "battle_outcome",
+                "prompt_name", "game_type", "board_size",
+                "outcome",
+            ],
+            quoting=csv.QUOTE_MINIMAL,
+            escapechar="\\"
+        )
+        self._writer.writeheader()
+        self._rows: List[dict] = []  # Store rows for post-game update
+
+    def log_move(
+        self,
+        turn: int,
+        player: int,
+        move: str,
+        model_name: str = "",
+        src: str = "",
+        dst: str = "",
+        piece_type: str = "",
+        outcome: str = "",
+        board_state: str = "",
+        available_moves: str = "",
+        move_direction: str = "",
+        target_piece: str = "",
+        battle_outcome: str = "",
+    ):
+        row = {
+            "turn": turn,
+            "player": player,
+            "model_name": model_name,
+            "move": move,
+            "from_pos": src,
+            "to_pos": dst,
+            "piece_type": piece_type,
+            "board_state": board_state,
+            "available_moves": available_moves,
+            "move_direction": move_direction,
+            "target_piece": target_piece,
+            "battle_outcome": battle_outcome,
+            "prompt_name": self.prompt_name,
+            "game_type": self.game_type,
+            "board_size": self.board_size,
+            "outcome": outcome,
+        }
+
+        # Store row in memory for post-game rewrite and write immediately to the file
+        self._rows.append(row)
+        try:
+            self._writer.writerow(row)
+            self._f.flush()
+        except Exception:
+            # If writing fails, keep going (file may be temporarily locked)
+            pass
+    
+    def finalize_game(self, winner: Optional[int], game_result: str = ""):
+        """
+        Rewrite CSV with game outcome in every row.
+        This allows each move to be labeled with the game result for training.
+        """
+        self._f.close()
+        
+        # Update all rows with game outcome
+        winner_str = str(winner) if winner is not None else "draw"
+        for row in self._rows:
+            row["game_winner"] = winner_str
+            row["game_result"] = game_result
+        
+        # Rewrite the file with updated data
+        with open(self.path, "w", newline="", encoding="utf-8") as f:
+            writer = csv.DictWriter(
+                f,
+                fieldnames=[
+                    "turn", "player", "model_name",
+                    "move", "from_pos", "to_pos", "piece_type",
+                    "board_state", "available_moves", "move_direction",
+                    "target_piece", "battle_outcome",
+                    "prompt_name", "game_type", "board_size",
+                    "outcome", "game_winner", "game_result",
+                ],
+                quoting=csv.QUOTE_MINIMAL,
+                escapechar="\\"
+            )
+            writer.writeheader()
+            writer.writerows(self._rows)
+    
+    def close(self):
+        try:
+            self._f.close()
+        except Exception:
+            pass
+    
+    def __enter__(self):
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+        return False