Spaces:
Running
Running
File size: 4,377 Bytes
3f9559b | 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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | """Human-like move prediction: CNN (piece + square classifiers) + Stockfish.
The CNN models predict the move a *human* would likely play (trained for white to
move; black positions are mirrored). Stockfish provides engine-strength moves and
a hybrid that keeps human-like moves which don't blunder.
"""
from __future__ import annotations
from pathlib import Path
import chess
import numpy as np
import onnxruntime as ort
PIECE_ORDER = [chess.PAWN, chess.KNIGHT, chess.BISHOP, chess.ROOK, chess.QUEEN, chess.KING]
SQUARE_NAMES = ["pawn", "knight", "bishop", "rook", "queen", "king"]
def _softmax(x):
e = np.exp(x - x.max())
return e / e.sum()
class MovePredictor:
def __init__(self, models_dir: str, stockfish_path: str | None = None):
d = Path(models_dir)
prov = ["CPUExecutionProvider"]
self.piece = ort.InferenceSession(str(d / "piece.int8.onnx"), providers=prov)
self.squares = [ort.InferenceSession(str(d / f"square_{s}.int8.onnx"), providers=prov)
for s in SQUARE_NAMES]
self.stockfish_path = (stockfish_path if stockfish_path and Path(stockfish_path).exists()
else None)
# ── encoding: 12x8x8, ch 0-5 white P,N,B,R,Q,K; 6-11 black; row0 = rank8 ──
def _encode(self, board: chess.Board):
enc = np.zeros((12, 8, 8), np.float32)
for sq in chess.SQUARES:
p = board.piece_at(sq)
if p:
c = PIECE_ORDER.index(p.piece_type) + (0 if p.color == chess.WHITE else 6)
enc[c, 7 - (sq >> 3), sq & 7] = 1.0
return enc[None]
def _cnn_white(self, board: chess.Board) -> dict:
enc = self._encode(board)
pieces = _softmax(self.piece.run(None, {"input": enc})[0].flatten())
legal = set(board.legal_moves)
mp: dict[str, float] = {}
for i, ptype in enumerate(PIECE_ORDER):
sq = (_softmax(self.squares[i].run(None, {"input": enc})[0].flatten()) + pieces[i]) / 2
for fsq in board.pieces(ptype, chess.WHITE):
fs = chess.square_name(fsq)
for j in range(64):
try:
mv = chess.Move.from_uci(fs + chess.square_name(j))
except ValueError:
continue
if mv in legal:
mp[mv.uci()] = float(sq[j])
return mp
def cnn_scores(self, board: chess.Board) -> dict:
"""Move->score; handles black by mirroring the board and the moves."""
if board.turn == chess.WHITE:
return self._cnn_white(board)
mirrored = self._cnn_white(board.mirror())
out = {}
for uci, s in mirrored.items():
mv = chess.Move.from_uci(uci)
real = chess.Move(chess.square_mirror(mv.from_square),
chess.square_mirror(mv.to_square), mv.promotion)
out[real.uci()] = s
return out
def predict(self, fen: str, top_n: int = 3) -> dict:
board = chess.Board(fen)
cnn = [m for m, _ in sorted(self.cnn_scores(board).items(), key=lambda x: -x[1])]
out = {"fen": fen, "turn": "white" if board.turn else "black",
"cnn": cnn[:top_n], "stockfish": [], "hybrid": cnn[:top_n]}
if self.stockfish_path:
from stockfish import Stockfish
sf = Stockfish(path=self.stockfish_path, parameters={"Threads": 1})
sf.set_fen_position(fen)
sf.update_engine_parameters({"MultiPV": top_n})
out["stockfish"] = [t["Move"] for t in sf.get_top_moves(top_n) if t.get("Move")]
# hybrid: among the top human-like CNN moves, prefer the one Stockfish
# rates best (lowest eval for the opponent after the move).
scored = []
for mv in cnn[:max(6, top_n * 2)]:
board.push_uci(mv)
sf.set_fen_position(board.fen())
ev = sf.get_evaluation()
val = ev.get("value", 0) if ev.get("type") == "cp" else (10000 if ev.get("value", 0) > 0 else -10000)
board.pop()
scored.append((mv, val)) # val = opponent's eval after our move (lower = better for us)
out["hybrid"] = [m for m, _ in sorted(scored, key=lambda x: x[1])][:top_n]
return out
|