Update chess-gpt-eval-contrastive/mamba_module.py

chess-gpt-eval-contrastive/mamba_module.py

@@ -6,6 +6,7 @@ from mamba_ssm import MambaLMHeadModel
 from contextlib import nullcontext
 import numpy as np
 from functools import partial
+import chess
 
 BASE_DIR = "mamba/"
 
@@ -176,7 +177,7 @@ class MambaPlayer:
         except:
             return None
 
-    def get_move(self, board: str, game_state: str, temperature: float) -> str:
+    def get_move(self, board: chess.Board, game_state: str, temperature: float) -> str:
         self.move_num = game_state.count('.')
         completion = self.get_mamba_response(game_state, temperature, 8, self.vocab_size)
         return self.get_move_from_response(completion)