rulesmith 0.1.0__py3-none-any.whl

rulesmith/chess.py

@@ -0,0 +1,526 @@
+"""Score decision graphs by playing chess, one graph execution per move.
+
+A graph cannot name a move: there are thousands and which are legal changes every ply. It names
+a policy instead -- take the best capture, get the king castled, save the piece that is hanging --
+and the board turns that into a move, deterministically. The labels are the policies, so the same
+rule language that presses buttons in Doom plays chess here.
+"""
+
+import random
+from typing import Literal, Self
+
+import chess
+import chess.engine
+from pydantic import Field, model_validator
+
+from rulesmith.graph import Plan, StrictModel, Text
+from rulesmith.optimize import Outcome
+from rulesmith.runtime import DecisionProgram, undecided
+
+VALUES = {chess.PAWN: 1, chess.KNIGHT: 3, chess.BISHOP: 3, chess.ROOK: 5, chess.QUEEN: 9}
+CENTER = [chess.D4, chess.E4, chess.D5, chess.E5]
+# Material left on both sides, below which the game is played as an endgame.
+ENDGAME = 24
+# The edge that decides a game stopped at the move limit: a minor piece, not a pawn.
+MARGIN = 3
+
+
+def worth(piece: chess.Piece | None) -> int:
+    return VALUES.get(piece.piece_type, 0) if piece else 0
+
+
+def gain(board: chess.Board, move: chess.Move) -> int:
+    """What a move wins in material, less the mover if the square can be taken back. An
+    approximation of an exchange search: one recapture deep, which is what the labels need."""
+    taken = worth(board.piece_at(move.to_square))
+    if board.is_en_passant(move):
+        taken = VALUES[chess.PAWN]
+    mover = worth(board.piece_at(move.from_square))
+    if move.promotion:
+        # The pawn becomes the new piece, which is both the gain and what can be taken back.
+        taken += VALUES[move.promotion] - VALUES[chess.PAWN]
+        mover = VALUES[move.promotion]
+    board.push(move)
+    exposed = board.is_attacked_by(board.turn, move.to_square)
+    board.pop()
+    return taken - (mover if exposed else 0)
+
+
+def centrality(square: int) -> int:
+    """How near the middle a square is; ties between developing moves break on this."""
+    file, rank = chess.square_file(square), chess.square_rank(square)
+    return -(abs(2 * file - 7) + abs(2 * rank - 7))
+
+
+def hanging(board: chess.Board) -> list[tuple[int, int]]:
+    """Own pieces the opponent attacks and this side does not defend, worst first."""
+    loose = []
+    for square, piece in board.piece_map().items():
+        if piece.color != board.turn or piece.piece_type == chess.KING:
+            continue
+        if board.is_attacked_by(not board.turn, square) and not board.is_attacked_by(
+            board.turn, square
+        ):
+            loose.append((worth(piece), square))
+    return sorted(loose, reverse=True)
+
+
+def quiet(board: chess.Board, move: chess.Move) -> bool:
+    return not board.is_capture(move) and not board.gives_check(move)
+
+
+def take(board: chess.Board) -> chess.Move | None:
+    captures = [move for move in board.legal_moves if board.is_capture(move)]
+    return max(captures, key=lambda move: (gain(board, move), move.uci()), default=None)
+
+
+def mates(board: chess.Board, move: chess.Move) -> bool:
+    board.push(move)
+    over = board.is_checkmate()
+    board.pop()
+    return over
+
+
+def check(board: chess.Board) -> chess.Move | None:
+    """A checking move, and if one of them ends the game, that one: a policy that says 'give
+    check' and passes over mate in one is not doing what it says."""
+    checks = [move for move in board.legal_moves if board.gives_check(move)]
+    return max(
+        checks, key=lambda move: (mates(board, move), gain(board, move), move.uci()), default=None
+    )
+
+
+def castle(board: chess.Board) -> chess.Move | None:
+    castles = [move for move in board.legal_moves if board.is_castling(move)]
+    return min(castles, key=lambda move: move.uci(), default=None)
+
+
+def develop(board: chess.Board) -> chess.Move | None:
+    """A knight or bishop off its own back rank, towards the middle."""
+    home = 0 if board.turn == chess.WHITE else 7
+    moves = [
+        move
+        for move in board.legal_moves
+        if (piece := board.piece_at(move.from_square))
+        and piece.piece_type in (chess.KNIGHT, chess.BISHOP)
+        and chess.square_rank(move.from_square) == home
+        and gain(board, move) >= 0
+    ]
+    return max(moves, key=lambda move: (centrality(move.to_square), move.uci()), default=None)
+
+
+def advance(board: chess.Board) -> chess.Move | None:
+    """A pawn forward: into the middle early, up the board once the pieces have gone."""
+    pawns = [
+        move
+        for move in board.legal_moves
+        if (piece := board.piece_at(move.from_square))
+        and piece.piece_type == chess.PAWN
+        and gain(board, move) >= 0
+    ]
+    ahead = 1 if board.turn == chess.WHITE else -1
+    return max(
+        pawns,
+        key=lambda move: (
+            move.to_square in CENTER,
+            ahead * (chess.square_rank(move.to_square) - chess.square_rank(move.from_square)),
+            centrality(move.to_square),
+            move.uci(),
+        ),
+        default=None,
+    )
+
+
+def save(board: chess.Board) -> chess.Move | None:
+    """Move the most valuable hanging piece somewhere it is not attacked."""
+    for _, square in hanging(board):
+        escapes = [
+            move
+            for move in board.legal_moves
+            if move.from_square == square and gain(board, move) >= 0
+        ]
+        safe = [move for move in escapes if not exposes(board, move)]
+        if safe or escapes:
+            return max(
+                safe or escapes, key=lambda move: (gain(board, move), centrality(move.to_square))
+            )
+    return None
+
+
+def exposes(board: chess.Board, move: chess.Move) -> bool:
+    board.push(move)
+    attacked = board.is_attacked_by(board.turn, move.to_square)
+    board.pop()
+    return attacked
+
+
+def settle(board: chess.Board) -> chess.Move:
+    """The move played when a policy has nothing to offer: a quiet move that keeps the piece it
+    moves, else the one that loses least. It must not reach for a capture the policy did not
+    ask for -- a fallback that grabs material makes every policy the same policy."""
+    moves = sorted(board.legal_moves, key=lambda move: move.uci())
+    calm = [move for move in moves if quiet(board, move) and not exposes(board, move)]
+    if calm:
+        return max(calm, key=lambda move: (centrality(move.to_square), move.uci()))
+    return max(moves, key=lambda move: (gain(board, move), move.uci()))
+
+
+def loose_for(board: chess.Board, color: bool) -> list[tuple[int, int]]:
+    """Pieces of one color the other side attacks and that color does not defend, worst first."""
+    found = []
+    for square, piece in board.piece_map().items():
+        if piece.color != color or piece.piece_type == chess.KING:
+            continue
+        if board.is_attacked_by(not color, square) and not board.is_attacked_by(color, square):
+            found.append((worth(piece), square))
+    return sorted(found, reverse=True)
+
+
+def ring(board: chess.Board, color: bool) -> int:
+    """How many squares around a king the other side attacks: pressure before it becomes check."""
+    king = board.king(color)
+    if king is None:
+        return 0
+    around = chess.SquareSet(chess.BB_KING_ATTACKS[king])
+    return sum(1 for square in around if board.is_attacked_by(not color, square))
+
+
+def passed(board: chess.Board, color: bool, square: int) -> bool:
+    """No enemy pawn ahead of this one on its file or the ones beside it."""
+    ahead = 1 if color == chess.WHITE else -1
+    file, rank = chess.square_file(square), chess.square_rank(square)
+    return not any(
+        chess.square_file(other) in (file - 1, file, file + 1)
+        and (chess.square_rank(other) - rank) * ahead > 0
+        for other in board.pieces(chess.PAWN, not color)
+    )
+
+
+def passers(board: chess.Board, color: bool) -> list[int]:
+    """Own passed pawns, by how far up the board they are."""
+    return sorted(
+        (
+            chess.square_rank(square) if color == chess.WHITE else 7 - chess.square_rank(square)
+            for square in board.pieces(chess.PAWN, color)
+            if passed(board, color, square)
+        ),
+        reverse=True,
+    )
+
+
+def castled(board: chess.Board, color: bool) -> bool:
+    """Whether this side castled, not merely moved its king. Replaying the game is the only
+    way to tell, so it is done only when the king stands where castling puts it."""
+    home = (chess.G1, chess.C1) if color == chess.WHITE else (chess.G8, chess.C8)
+    if board.king(color) not in home:
+        return False
+    replay = board.root()
+    for move in board.move_stack:
+        if replay.turn == color and replay.is_castling(move):
+            return True
+        replay.push(move)
+    return False
+
+
+def promote(board: chess.Board) -> chess.Move | None:
+    """Push the passed pawn that is furthest up the board, promoting when it can."""
+    ahead = 1 if board.turn == chess.WHITE else -1
+    free = {
+        square
+        for square in board.pieces(chess.PAWN, board.turn)
+        if passed(board, board.turn, square)
+    }
+    pushes = [
+        move
+        for move in board.legal_moves
+        if move.from_square in free
+        and (piece := board.piece_at(move.from_square))
+        and piece.piece_type == chess.PAWN
+        and (chess.square_rank(move.to_square) - chess.square_rank(move.from_square)) * ahead > 0
+        and gain(board, move) >= 0
+    ]
+    if not pushes:
+        return None
+
+    def reach(move: chess.Move) -> int:
+        rank = chess.square_rank(move.to_square)
+        return rank if board.turn == chess.WHITE else 7 - rank
+
+    return max(pushes, key=lambda move: (move.promotion or 0, reach(move), move.uci()))
+
+
+POLICIES = {
+    "CAPTURE": take,
+    "CHECK": check,
+    "CASTLE": castle,
+    "DEVELOP": develop,
+    "ADVANCE": advance,
+    "SAVE": save,
+    "PROMOTE": promote,
+    "QUIET": lambda board: None,
+}
+
+
+def chosen(board: chess.Board, label: str) -> chess.Move:
+    """The move a policy names in this position; every policy always yields a legal move."""
+    if label not in POLICIES:
+        raise ValueError(f"unknown policy {label!r}; the task allows {sorted(POLICIES)}")
+    return POLICIES[label](board) or settle(board)
+
+
+def observe(board: chess.Board) -> dict:
+    """What the side to move can see, as numbers a read node can take by path."""
+    mine = sum(worth(piece) for piece in board.piece_map().values() if piece.color == board.turn)
+    theirs = sum(worth(piece) for piece in board.piece_map().values() if piece.color != board.turn)
+    loose = hanging(board)
+    theirs_loose = loose_for(board, not board.turn)
+    mine_passed = passers(board, board.turn)
+    best = take(board)
+    checking = check(board)
+    back = 0 if board.turn == chess.WHITE else 7
+    return {
+        "ply": board.ply(),
+        "phase": "endgame"
+        if mine + theirs <= ENDGAME
+        else "opening"
+        if board.ply() < 20
+        else "middlegame",
+        "material": {"mine": mine, "theirs": theirs, "balance": mine - theirs},
+        "king": {
+            "in_check": int(board.is_check()),
+            "attackers": len(board.attackers(not board.turn, board.king(board.turn))),
+            "can_castle": int(bool(castle(board))),
+            "castled": int(castled(board, board.turn)),
+        },
+        "capture": {
+            "best_gain": gain(board, best) if best else 0,
+            "count": sum(1 for move in board.legal_moves if board.is_capture(move)),
+        },
+        "check": {
+            "available": int(bool(checking)),
+            "gain": gain(board, checking) if checking else 0,
+        },
+        "hanging": {"count": len(loose), "worst": loose[0][0] if loose else 0},
+        # What is loose on the other side is a reason to go looking, not only a capture to take.
+        "target": {
+            "loose_count": len(theirs_loose),
+            "loose_worst": theirs_loose[0][0] if theirs_loose else 0,
+            "king_ring": ring(board, not board.turn),
+        },
+        "pawns": {
+            "passed": len(mine_passed),
+            "furthest": mine_passed[0] if mine_passed else 0,
+            "theirs_passed": len(passers(board, not board.turn)),
+        },
+        "safety": {"king_ring": ring(board, board.turn)},
+        "development": {
+            "minor_home": sum(
+                1
+                for square, piece in board.piece_map().items()
+                if piece.color == board.turn
+                and piece.piece_type in (chess.KNIGHT, chess.BISHOP)
+                and chess.square_rank(square) == back
+            )
+        },
+        "moves": {"count": board.legal_moves.count()},
+    }
+
+
+class ChessTask(StrictModel):
+    game: Literal["chess"] = "chess"
+    description: Text
+    labels: list[Text] = Field(min_length=2, description="Policies the graph may name")
+    opening_plies: int = Field(default=4, ge=0, description="Random legal moves opening a game")
+    move_limit: int = Field(default=120, ge=2, description="Plies before the game is adjudicated")
+    material_weight: float = Field(
+        default=0.2,
+        ge=0,
+        le=0.5,
+        description="Share of the score decided by material rather than by the result, so that "
+        "a game nobody wins still says which graph stood better",
+    )
+    survival_weight: float = Field(
+        default=0,
+        ge=0,
+        le=0.5,
+        description="Share of the score decided by how long a lost or drawn game lasted, so "
+        "that against an opponent no graph can beat, holding out longer still scores higher",
+    )
+    train: list[int] = Field(min_length=1)
+    validation: list[int] = Field(min_length=1)
+    test: list[int] = Field(min_length=1)
+
+    @model_validator(mode="after")
+    def consistent(self) -> Self:
+        if len(set(self.labels)) != len(self.labels):
+            raise ValueError("labels must be distinct")
+        unknown = set(self.labels) - set(POLICIES)
+        if unknown:
+            raise ValueError(f"unknown policies: {sorted(unknown)}; allowed {sorted(POLICIES)}")
+        seeds = self.train + self.validation + self.test
+        if len(set(seeds)) != len(seeds):
+            raise ValueError("seeds must be distinct across splits")
+        # A win has to outrank any other result whatever the material, which holds while the
+        # result's share of the score is more than twice material's.
+        if 1 - 3 * self.material_weight - self.survival_weight <= 0:
+            raise ValueError("material and survival weights leave the result too little weight")
+        return self
+
+    def baseline_plan(self) -> Plan:
+        return Plan.constant(self.labels[0])
+
+    def opening(self, seed: int) -> chess.Board:
+        """A board a few random legal moves in, so that one seed is one game."""
+        board = chess.Board()
+        rng = random.Random(seed)
+        for _ in range(self.opening_plies):
+            moves = sorted(board.legal_moves, key=lambda move: move.uci())
+            if not moves or board.is_game_over():
+                break
+            board.push(rng.choice(moves))
+        return board
+
+
+class Engine:
+    """A UCI engine at a fixed strength, as a reference opponent that does not drift the way
+    Elo fitted over a pool of graphs does. Stockfish's skill level 0 is calibrated near 1320
+    CCRL blitz Elo."""
+
+    def __init__(self, path: str, skill: int, seconds: float):
+        self.path, self.skill, self.limit = path, skill, chess.engine.Limit(time=seconds)
+
+    def __enter__(self) -> Self:
+        self.engine = chess.engine.SimpleEngine.popen_uci(self.path)
+        self.engine.configure({"Skill Level": self.skill})
+        return self
+
+    def __exit__(self, *raised) -> None:
+        self.engine.quit()
+
+    def move(self, board: chess.Board) -> chess.Move:
+        return self.engine.play(board, self.limit).move
+
+
+Player = DecisionProgram | Engine
+
+
+def played(task: ChessTask, white: Player, black: Player, seed: int) -> dict:
+    """One game. Returns the result from white's side, with what each side did along the way."""
+    board = task.opening(seed)
+    sides = {chess.WHITE: white, chess.BLACK: black}
+    moves, models, error, loser = [], set(), None, None
+    first: dict[bool, dict] = {}
+    while not board.is_game_over() and board.ply() < task.move_limit:
+        program = sides[board.turn]
+        if isinstance(program, Engine):
+            move = program.move(board)
+            moves.append({"policy": "engine", "move": move.uci()})
+            board.push(move)
+            continue
+        seen = observe(board)
+        first.setdefault(board.turn, seen)
+        prediction = program(state=seen)
+        models.update(prediction.models)
+        if prediction.undecided:
+            error, loser = undecided(program.plan, prediction), board.turn
+            break
+        try:
+            move = chosen(board, prediction.label)
+        except ValueError as failure:
+            error, loser = str(failure), board.turn
+            break
+        moves.append({"policy": prediction.label, "move": move.uci()})
+        board.push(move)
+    edge = material(board, chess.WHITE)
+    if board.is_game_over():
+        outcome = board.result(claim_draw=True)
+    else:
+        # Two graphs that cannot mate will shuffle to the move limit, and calling every one of
+        # those a draw leaves search nothing to climb. A game stopped a piece down was lost.
+        outcome = "1-0" if edge >= MARGIN else "0-1" if edge <= -MARGIN else "1/2-1/2"
+    if loser is not None:
+        outcome = "0-1" if loser == chess.WHITE else "1-0"
+    return {
+        "result": outcome,
+        "moves": moves,
+        "fen": board.fen(),
+        "material": edge,
+        "adjudicated": not board.is_game_over() and loser is None,
+        "models": sorted(models),
+        "error": error,
+        "failed": loser,
+        # What each side saw on its first move, which is where its choices began.
+        "first": first,
+    }
+
+
+def material(board: chess.Board, color: bool) -> int:
+    ours = sum(worth(p) for p in board.piece_map().values() if p.color == color)
+    theirs = sum(worth(p) for p in board.piece_map().values() if p.color != color)
+    return ours - theirs
+
+
+def graded(task: ChessTask, points: float, edge: int, plies: int) -> float:
+    """A game's score: mostly its result, partly the material left, and partly how long it
+    lasted. A win counts as lasting the whole game, since ending it sooner is no worse."""
+    material, survival = task.material_weight, task.survival_weight
+    held = 1.0 if points == 1 else min(1.0, plies / (task.move_limit - task.opening_plies))
+    return (
+        (1 - material - survival) * points
+        + material * (0.5 + max(-15, min(15, edge)) / 30)
+        + survival * held
+    )
+
+
+class Games:
+    """Scores a graph by how it fares against a fixed opponent, one game per seed."""
+
+    def __init__(self, task: ChessTask, opponent: Player):
+        self.task = task
+        self.opponent = opponent
+
+    def __call__(self, program: DecisionProgram, seed: int) -> Outcome:
+        # Colors alternate by seed, so across a split a graph plays as many games as each side.
+        mine = chess.WHITE if seed % 2 == 0 else chess.BLACK
+        white, black = (program, self.opponent) if mine == chess.WHITE else (self.opponent, program)
+        game = played(self.task, white, black, seed)
+        points = {"1-0": 1.0, "0-1": 0.0, "1/2-1/2": 0.5, "*": 0.5}[game["result"]]
+        if mine == chess.BLACK:
+            points = 1 - points
+        edge = game["material"] if mine == chess.WHITE else -game["material"]
+        score = graded(self.task, points, edge, len(game["moves"]))
+        said = "won" if points == 1 else "lost" if points == 0 else "drew"
+        # Only this graph's own failure is its error; an opponent that breaks simply loses.
+        error = game["error"] if game["failed"] == mine else None
+        if error is None and game["error"] is not None:
+            said = f"won because the opponent failed ({game['error']})"
+        return Outcome(
+            score=score if error is None else 0.0,
+            record={
+                "seed": seed,
+                "color": "white" if mine == chess.WHITE else "black",
+                "result": game["result"],
+                "points": points,
+                "material": edge,
+                "plies": len(game["moves"]),
+                "policies": {
+                    policy: sum(1 for m in game["moves"] if m["policy"] == policy)
+                    for policy in sorted({m["policy"] for m in game["moves"]})
+                },
+                "score": score,
+                "error": error,
+                "fen": game["fen"],
+            },
+            trace={
+                "Inputs": game["first"].get(mine, {}),
+                "Generated Outputs": {"moves": game["moves"][-8:]},
+                "Feedback": error
+                or (
+                    f"Playing {'white' if mine == chess.WHITE else 'black'} from seed {seed} you "
+                    f"{said} in {len(game['moves'])} plies, ending {edge:+d} in material "
+                    f"(score {score:.3f}). Higher is better."
+                ),
+            },
+            models=game["models"],
+            error=error,
+        )

rulesmith/clef.py

@@ -0,0 +1,66 @@
+"""Clef, Cloudflare's decision model, served from its own release: `POST /v1/systemone`
+answered as Jev answers it, on an NVIDIA GPU, Apple's, or a CPU, for `--backend clef`."""
+
+import importlib
+import sys
+import threading
+import traceback
+from collections.abc import Callable
+from http import HTTPStatus
+from http.server import ThreadingHTTPServer
+from pathlib import Path
+
+from rulesmith.serve import JSONHandler
+
+PATH = "/v1/systemone"
+
+Answer = Callable[[dict], dict]
+
+
+def load(weights: str, device: str | None) -> tuple[Answer, str]:
+    """The release's own `systemone` bound to its model, and the device it runs on: `device`,
+    or else the best this machine has. `weights` is a Hugging Face repo or a downloaded copy."""
+    # The clef extra's packages, imported only here so the rest of rulesmith installs without.
+    import torch
+    from huggingface_hub import snapshot_download
+
+    if device is None:
+        if torch.cuda.is_available():
+            device = "cuda"
+        elif torch.backends.mps.is_available():
+            device = "mps"
+        else:
+            device = "cpu"
+            print("warning: no GPU found; Clef will run on the CPU, slowly", file=sys.stderr)
+    path = Path(weights) if Path(weights).is_dir() else Path(snapshot_download(weights))
+    # The model's code ships beside its weights, as the release's model card loads it.
+    sys.path.insert(0, str(path))
+    release = importlib.import_module("joint_schema_model")
+    model, processor = release.load_release_model(path, device=device)
+    return (lambda request: release.systemone(model, processor, request)), device
+
+
+def make_server(answer: Answer, host: str, port: int) -> ThreadingHTTPServer:
+    """A server that answers each systemone request with `answer`, one at a time."""
+
+    class Handler(JSONHandler):
+        def post(self, request):
+            if self.path != PATH:
+                return HTTPStatus.NOT_FOUND, {"error": f"nothing at {self.path}; ask {PATH}"}
+            if not isinstance(request, dict):
+                return HTTPStatus.BAD_REQUEST, {"error": "a systemone request is a JSON object"}
+            try:
+                with self.server.lock:
+                    return HTTPStatus.OK, answer(request)
+            # The release checks a request as it reads it, so a malformed one fails here.
+            except (ValueError, KeyError, TypeError, AttributeError) as error:
+                return HTTPStatus.BAD_REQUEST, {"error": f"not a systemone request: {error}"}
+            except RuntimeError as error:
+                # Out of memory, most often: the caller hears why, the log keeps where.
+                traceback.print_exc()
+                return HTTPStatus.INTERNAL_SERVER_ERROR, {"error": f"the model failed: {error}"}
+
+    server = ThreadingHTTPServer((host, port), Handler)
+    # One model on one device, which the release does not make safe to run from two threads.
+    server.lock = threading.Lock()
+    return server