| from __future__ import annotations |
| import os |
| import random |
| import sys |
| import chess |
| from search import Searcher |
|
|
| ENGINE_NAME = "ChessMamba" |
| ENGINE_AUTHOR = "CatGirlHtfi" |
| DEFAULT_CKPT = os.environ.get( |
| "CHESSMAMBA_CKPT", os.path.join(os.path.dirname(__file__), "ckpt", "model.pt") |
| ) |
| POLICY_ONLY = os.environ.get("CHESSMAMBA_POLICY_ONLY", "") == "1" |
|
|
|
|
| def send(msg: str) -> None: |
| sys.stdout.write(msg + "\n") |
| sys.stdout.flush() |
|
|
|
|
| def main() -> None: |
| board = chess.Board() |
| searcher = Searcher(checkpoint_path=DEFAULT_CKPT, policy_only=POLICY_ONLY) |
| for line in sys.stdin: |
| line = line.strip() |
| if not line: |
| continue |
| tokens = line.split() |
| cmd = tokens[0] |
| if cmd == "uci": |
| send(f"id name {ENGINE_NAME}") |
| send(f"id author {ENGINE_AUTHOR}") |
| send("uciok") |
| elif cmd == "isready": |
| send("readyok") |
| elif cmd == "ucinewgame": |
| board = chess.Board() |
| searcher.reset() |
| elif cmd == "position": |
| _handle_position(board, tokens) |
| searcher.sync(board) |
| elif cmd == "go": |
| move = searcher.choose_move(board, tokens) |
| if move is None: |
| legal = list(board.legal_moves) |
| move = random.choice(legal) if legal else None |
| if move is not None: |
| send(f"bestmove {move.uci()}") |
| else: |
| send("bestmove 0000") |
| elif cmd == "quit": |
| break |
| elif cmd == "stop": |
| pass |
|
|
|
|
| def _handle_position(board: chess.Board, tokens: list[str]) -> None: |
| if "startpos" in tokens: |
| board.reset() |
| idx = tokens.index("startpos") + 1 |
| elif "fen" in tokens: |
| fen_idx = tokens.index("fen") |
| moves_idx = tokens.index("moves") if "moves" in tokens else len(tokens) |
| fen = " ".join(tokens[fen_idx + 1 : moves_idx]) |
| board.set_fen(fen) |
| idx = moves_idx |
| else: |
| return |
| if idx < len(tokens) and tokens[idx] == "moves": |
| for uci_move in tokens[idx + 1 :]: |
| board.push_uci(uci_move) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
