File size: 2,183 Bytes
c0c6c62 380c43e c0c6c62 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | 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()
|