""" complete_deck.py CLI: given a commander (or two, for partners) and an optional partial decklist by name, print the model's top-K suggested completions -- the actual deliverable of the project, not just an evaluation number. Applies the deterministic color-identity legality filter described in the training plan (a candidate's color identity must be a subset of the commander(s)' color identity) as a post-filter on the model's raw scores, rather than relying on the model to have learned the rule from data. Generation strategy: single-shot (default) vs. iterative (--chunk-size / --temperature) -- see complete_deck()'s docstring below for what these do and why they exist. Single-shot (chunk_size=None) is unchanged from the original implementation, so every existing example in this file's docstring and README.md still reproduces exactly. Usage: python3 -m training.complete_deck --commander "Atraxa, Praetors' Voice" python3 -m training.complete_deck --commander "Atraxa, Praetors' Voice" \\ --card "Sol Ring" --card "Arcane Signet" --top 15 python3 -m training.complete_deck \\ --commander "Silas Renn, Seeker Adept" --commander "Rograkh, Son of Rohgahh" python3 -m training.complete_deck --commander "Meren of Clan Nel Toth" \\ --chunk-size 5 --temperature 0.5 """ from __future__ import annotations import argparse import random from collections.abc import Iterable import numpy as np import torch from tokenizer.mtg_tokenizer import COLORS, SPECIAL_TOKENS, CardTokenizer from training.evaluate import DEFAULT_CHECKPOINT, load_checkpoint from training.model import DeckCompletionModel from training.train import pick_device class UnknownCardError(ValueError): """Raised when a given commander/card name doesn't resolve against the tokenizer's vocab. A plain ValueError subclass (not SystemExit) so non-CLI callers can catch it and turn it into a proper error response instead of killing the process; __main__ below is the one place that still turns it into a CLI exit.""" def legal_mask( tok: CardTokenizer, commander_colors: set[str], allow_banned: bool = True, exclude_ids: Iterable[int] = (), vocab_size: int | None = None, ) -> torch.Tensor: """(vocab_size,) bool tensor: True where a card's color identity is a subset of commander_colors. Special tokens (all-zero feature rows) are excluded explicitly, since an empty color identity would otherwise trivially pass the subset check. allow_banned=False also excludes every card whose Commander legality is "banned" as of the tokenizer's Scryfall snapshot -- the vocab keeps banned cards on purpose (see CardRecord.commander_legality), so the model can and does suggest them unless a caller filters them out here. exclude_ids: token ids to drop from the candidate pool regardless of legality -- e.g. cards a user already rejected, so they don't come back. Ids past vocab_size (below) are outside the mask already and are skipped. vocab_size: truncates to this many rows instead of tok's own (current) vocab_size -- pass the specific model's card_embedding.vocab_size at inference time. Matters for the same reason resolve_or_die's max_token_id does: tokenizer.json can grow past a checkpoint's trained vocab (new sets get appended via previous_tokenizer_path -- see resolve_or_die's docstring), and this mask gets applied directly against that model's (fixed-size) logits via logits[~mask] -- a mask sized to tok's current, possibly-larger vocab_size would mismatch shape against the model's actual (smaller, frozen-at-training-time) output. New cards are always appended at the end, so truncating to the first vocab_size rows never misclassifies any card the model actually knows about.""" feats = tok.feature_matrix() if vocab_size is not None: feats = feats[:vocab_size] color_cols = feats[:, : len(COLORS)] # (vocab_size, 5), 0/1 allowed = np.array([1.0 if c in commander_colors else 0.0 for c in COLORS], dtype=np.float32) illegal = (color_cols * (1.0 - allowed)).sum(axis=1) > 0 mask = torch.tensor(~illegal) for tok_str in SPECIAL_TOKENS: mask[tok.token_to_id[tok_str]] = False if not allow_banned: for tok_str, card in tok.card_by_token.items(): cid = tok.token_to_id[tok_str] if card.commander_legality == "banned" and cid < len(mask): mask[cid] = False for cid in exclude_ids: if cid < len(mask): mask[cid] = False return mask def resolve_or_die(tok: CardTokenizer, name: str, max_token_id: int | None = None) -> int: """Resolves name -> token id, then (if max_token_id is given) checks the id is actually within a loaded checkpoint's trained vocabulary before handing it back. This matters because the tokenizer and a checkpoint are independently-versioned artifacts -- tokenizer.json gets refreshed periodically as Scryfall prints new sets (fetch fresh oracle-cards.jsonl, then CardTokenizer.from_scryfall_bulk(..., previous_tokenizer_path=...) so existing cards keep their ids), appending new cards at the end of the vocab, but a checkpoint's HybridCardEmbedding.id_embedding table is a fixed size, frozen at whatever vocab_size existed when it was trained. A name that only resolves to one of the newly-appended cards passes tok.resolve_name() fine (it's a real, valid vocab entry) but is out of bounds for that specific checkpoint's embedding table -- verified live that this isn't just a theoretical concern: on MPS, indexing id_embedding out of bounds doesn't even raise a catchable Python exception, it aborts the whole process later (possibly after silently computing on garbage first). Catching it here, with the name still in hand, turns that into a clear, immediate UnknownCardError instead.""" token = tok.resolve_name(name) if token is None: raise UnknownCardError(f"Could not resolve card name: {name!r}") token_id = tok.token_to_id[token] if max_token_id is not None and token_id > max_token_id: raise UnknownCardError( f"{name!r} was added to the tokenizer after this checkpoint was trained " f"(token id {token_id}, checkpoint vocab covers only 0-{max_token_id}) -- " f"not usable with this model until it's retrained on a tokenizer that includes it." ) return token_id def _run_model(model, tok, device, cmdr_tensor, cmdr_mask, picked: list[int], mask: torch.Tensor) -> torch.Tensor: """One forward pass, masked to -inf for already-picked cards and illegal colors. Shared by both generation modes below.""" if picked: ctx_tensor = torch.tensor([picked], dtype=torch.long, device=device) ctx_mask = torch.ones_like(ctx_tensor, dtype=torch.bool) else: # Empty context is a valid case (commander-only completion) -- the # DeepSets encoder's mask-aware pool handles an all-False mask by # clamping the count to 1, giving a zero-vector context contribution # rather than dividing by zero. See training/model.py. ctx_tensor = torch.zeros((1, 1), dtype=torch.long, device=device) ctx_mask = torch.zeros((1, 1), dtype=torch.bool, device=device) logits = model(ctx_tensor, ctx_mask, cmdr_tensor, cmdr_mask)[0].cpu() already_visible = set(picked) | set(cmdr_tensor[0].tolist()) logits[list(already_visible)] = float("-inf") logits[~mask] = float("-inf") return logits def _choose_dynamic_chunk_size( logits: torch.Tensor, full_table: torch.Tensor, min_chunk: int, max_chunk: int, ) -> tuple[int, float | None]: """Maps this round's candidate diversity to a chunk size in [min_chunk, max_chunk], for chunk_size="dynamic" mode (see complete_deck()'s docstring). Diversity is 1 - mean pairwise cosine similarity among the top max_chunk legal/unpicked candidates' hybrid embeddings (full_table = model.card_embedding.full_table(), the same tied embedding table already used for scoring). Direction: low diversity (the model's current top picks already look like "more of the same") -> smaller chunk, so generation reacts and re-conditions sooner instead of locking in several near-duplicate picks at once -- this is a direct measurement of the actual self-reinforcing collapse mechanism documented for fixed small chunk sizes (each pick nudges the pooled deck vector further toward whatever it's already leaning on). High diversity (the top candidates span clearly different cards/roles) -> bigger chunk, since there's no visible redundancy risk to react against yet. Two score-magnitude-based signals were tried first and rejected before this one: softmax over the whole legal/unpicked pool (tens of thousands of candidates) never moved off min_chunk, since its max is bounded far below 1.0 by sheer pool size regardless of how peaked the model's preference is; sigmoid(top logit) saturated near 1.0 from round 1 on every commander tried, since an obviously-correct top pick (e.g. Sol Ring) is confidently scored independent of how far into generation we are -- it can't tell "just reinforced itself into a corner" apart from "this one card is just genuinely great," which is exactly the distinction this function needs to make. Returns (chunk_size, mean_cosine) -- mean_cosine is None (and chunk_size is min_chunk) when fewer than 2 legal/unpicked candidates remain; it's returned for logging/diagnostic use only, not used by complete_deck() itself.""" pool_ids = torch.nonzero(logits > float("-inf"), as_tuple=True)[0] k_probe = min(max_chunk, len(pool_ids)) if k_probe < 2: return min_chunk, None top_idx = torch.topk(logits[pool_ids], k_probe).indices cand_vecs = full_table[pool_ids[top_idx]] cand_vecs = cand_vecs / cand_vecs.norm(dim=1, keepdim=True).clamp_min(1e-8) sim = cand_vecs @ cand_vecs.T n = sim.shape[0] mean_cos = ((sim.sum() - sim.diagonal().sum()) / (n * (n - 1))).item() diversity = max(0.0, min(1.0, 1.0 - mean_cos)) chunk = min_chunk + diversity * (max_chunk - min_chunk) return max(min_chunk, min(max_chunk, round(chunk))), mean_cos def _select_chunk(logits: torch.Tensor, k: int, temperature: float | None, rng: random.Random) -> list[int]: """Selects up to k token ids from one set of (already-masked) logits. temperature=None is deterministic top-k (what the CLI's default and every single-shot call use). Otherwise samples k ids without replacement from softmax(logits/temperature) via sequential renormalization -- k is always small (<=~10 in practice) so the O(k^2) cost here is negligible.""" pool_ids = torch.nonzero(logits > float("-inf"), as_tuple=True)[0] k_eff = min(k, len(pool_ids)) if k_eff == 0: return [] if temperature is None: top_idx = torch.topk(logits[pool_ids], k_eff).indices return pool_ids[top_idx].tolist() remaining_ids = pool_ids.tolist() remaining_scores = logits[pool_ids].clone() picked = [] for _ in range(k_eff): probs = torch.softmax(remaining_scores / temperature, dim=0) pos = rng.choices(range(len(remaining_ids)), weights=probs.tolist(), k=1)[0] picked.append(remaining_ids[pos]) del remaining_ids[pos] remaining_scores = torch.cat([remaining_scores[:pos], remaining_scores[pos + 1:]]) return picked @torch.no_grad() def complete_deck( commanders: list[str], cards: list[str] | None = None, top: int = 20, checkpoint_path: str = DEFAULT_CHECKPOINT, loaded: tuple[DeckCompletionModel, CardTokenizer, torch.device] | None = None, chunk_size: int | str | None = None, temperature: float | None = None, seed: int = 0, allow_banned: bool = True, exclude_ids: Iterable[int] = (), chunk_size_range: tuple[int, int] = (2, 10), ) -> list[tuple[str, float]]: """`loaded`, if given, skips the (slow: disk + model init) checkpoint load and reuses an already-resident (model, tokenizer, device) -- for a caller that loads once at startup and must not re-load per request. The CLI path below leaves `loaded` unset and keeps loading fresh, as before. `chunk_size` / `temperature` select the generation strategy -- added in response to a real report that completions "feel like a card recommender, not a deck builder": individually reasonable-looking cards that don't add up to one coherent gameplan. Root cause: the default behavior below (chunk_size=None) runs ONE forward pass against the original `cards` and takes the top-`top` scores all at once, so pick #47 has no idea picks #1-46 exist -- nothing before this ever conditioned a new pick on the model's OWN prior picks in the same call, only on the caller's original input. chunk_size, when set, re-runs the forward pass every `chunk_size` picks instead of once for everything, feeding each new batch's picks back into the context first -- so pick #47 (if chunk_size divides evenly) really can react to #1-46. Small chunk sizes are maximally reactive but risk a self-reinforcing collapse: each pick nudges the deck-so-far representation further toward whatever it's already leaning on, making the next pick even more likely to be "more of the same" (real example found while testing this: a commander whose completion degenerated into a run of near-identical cheap filler spells under chunk_size=1). temperature counteracts that by sampling each chunk's picks from softmax(logits/temperature) instead of always taking the single top-scoring card, so a non-top pick occasionally wins and the run doesn't narrow into one repetitive lane. temperature is ignored when chunk_size is None (single-shot has no "chunk" to sample within). `allow_banned=False` excludes Commander-banned cards from the candidate pool, and `exclude_ids` drops specific cards (see legal_mask); the defaults keep the original behavior. Recommended starting point if you want to try iterative generation: chunk_size=5, temperature=0.5 -- the best compromise found across an 11-commander comparison (redundancy pulled back down from chunk_size=1's self-reinforcing collapse, while still reproducing chunk_size=1's clear win on at least one commander where single-shot stayed generic). It's a real, reportable improvement on one axis, not a settled best config, so don't take chunk_size=5/temperature=0.5 as more final than that. chunk_size="dynamic" (string, not an int) picks a chunk size every round instead of using one fixed value -- see _choose_dynamic_chunk_size()'s docstring for the diversity signal it's based on and the reasoning behind its direction. chunk_size_range=(min, max) bounds it (default (2, 10), matching the 5/10 fixed sizes already tested; the floor of 2 deliberately avoids ever reproducing chunk_size=1's collapse mode, and bounds the worst-case number of rounds/forward-passes per completion). Experimental, evaluated separately from the chunk_size=5 recommendation above -- not yet a replacement for it. """ if loaded is not None: model, tok, device = loaded else: device = pick_device() model, tok, _ckpt = load_checkpoint(checkpoint_path, device) max_token_id = model.card_embedding.vocab_size - 1 cmdr_ids = [resolve_or_die(tok, name, max_token_id) for name in commanders] card_ids = [resolve_or_die(tok, name, max_token_id) for name in (cards or [])] cmdr_tensor = torch.tensor([cmdr_ids], dtype=torch.long, device=device) cmdr_mask = torch.ones_like(cmdr_tensor, dtype=torch.bool) commander_colors: set[str] = set() for cid in cmdr_ids: card = tok.card_by_token[tok.id_to_token[cid]] commander_colors |= set(card.color_identity) mask = legal_mask(tok, commander_colors, allow_banned=allow_banned, exclude_ids=exclude_ids, vocab_size=max_token_id + 1) if chunk_size is None: # Single-shot: exactly the original implementation, so every # existing documented example reproduces byte-identically. logits = _run_model(model, tok, device, cmdr_tensor, cmdr_mask, card_ids, mask) scores = torch.sigmoid(logits) top_ids = torch.topk(scores, top).indices.tolist() return [(tok.card_by_token[tok.id_to_token[i]].name, scores[i].item()) for i in top_ids] rng = random.Random(seed) picked = list(card_ids) results: list[tuple[int, float]] = [] full_table = model.card_embedding.full_table() if chunk_size == "dynamic" else None while len(results) < top: remaining = top - len(results) logits = _run_model(model, tok, device, cmdr_tensor, cmdr_mask, picked, mask) if chunk_size == "dynamic": k_raw, _mean_cos = _choose_dynamic_chunk_size(logits, full_table, *chunk_size_range) k = min(k_raw, remaining) else: k = min(chunk_size, remaining) chosen = _select_chunk(logits, k, temperature, rng) if not chosen: break # ran out of legal, unpicked candidates for cid in chosen: picked.append(cid) results.append((cid, torch.sigmoid(logits[cid]).item())) return [(tok.card_by_token[tok.id_to_token[cid]].name, score) for cid, score in results] def _print_completion( commanders: list[str], cards: list[str], results: list[tuple[str, float]], chunk_size: int | str | None, temperature: float | None, ) -> None: cmdr_str = " + ".join(commanders) print(f"\nCommander(s): {cmdr_str}") if cards: print(f"Known cards ({len(cards)}): {', '.join(cards)}") else: print("Known cards: (none -- completing from commander alone)") if chunk_size is None: print("Generation strategy: single-shot") else: print(f"Generation strategy: iterative (chunk_size={chunk_size}, temperature={temperature})") print(f"\nTop {len(results)} suggestions:") for name, score in results: print(f" {score:.3f} {name}") def _chunk_size_arg(value: str) -> int | str: """argparse type for --chunk-size: an int, or the literal 'dynamic'.""" if value == "dynamic": return value try: return int(value) except ValueError: raise argparse.ArgumentTypeError(f"--chunk-size must be an int or 'dynamic', got {value!r}") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--commander", action="append", required=True, dest="commanders", help="Commander name; repeat for partner commanders.") parser.add_argument("--card", action="append", default=[], dest="cards", help="A card already in the deck; repeat for each known card.") parser.add_argument("--top", type=int, default=20) parser.add_argument("--checkpoint", default=DEFAULT_CHECKPOINT) parser.add_argument( "--chunk-size", type=_chunk_size_arg, default=None, help=( "Re-run the model every N picks instead of once for the whole completion " "(the default, single-shot). Each new batch of picks is conditioned on " "everything picked so far in this run -- including the model's own prior " "picks, not just --card -- fixing single-shot's tendency to produce a pile " "of independently-good-looking cards instead of one coherent gameplan. " "Small values react fastest but risk a self-reinforcing narrow collapse " "(see --temperature). Recommended starting point: 5. Pass the literal " "string 'dynamic' instead of a number to pick a chunk size every round " "from the diversity of the model's top candidates (see --chunk-size-range " "and complete_deck()'s docstring)." ), ) parser.add_argument( "--chunk-size-range", type=int, nargs=2, default=(2, 10), metavar=("MIN", "MAX"), help="Only used with --chunk-size dynamic: bounds on the per-round chunk size. Default: 2 10.", ) parser.add_argument( "--temperature", type=float, default=None, help=( "Only used with --chunk-size. Samples each chunk's picks from " "softmax(logit/temperature) instead of always taking the single top-scoring " "card, so a non-top pick occasionally wins instead of --chunk-size narrowing " "into one repetitive lane. Omit for deterministic greedy picks within each " "chunk. Recommended starting point: 0.5." ), ) args = parser.parse_args() try: results = complete_deck( args.commanders, args.cards, top=args.top, checkpoint_path=args.checkpoint, chunk_size=args.chunk_size, temperature=args.temperature, chunk_size_range=tuple(args.chunk_size_range), ) except UnknownCardError as e: raise SystemExit(str(e)) _print_completion(args.commanders, args.cards, results, args.chunk_size, args.temperature)