Spaces:
Running
Running
Instruct mode: chat models reply with the move they pick (best minus worst scoring, multi-turn)
8ee889d verified | """Instruct mode: chat models choose every move by replying with it. | |
| The model never sees the grid. Every legal placement is described by its outcome in a few words (lines, holes, stack | |
| height, surface), placements with the same description become one move, and the model gets up to GROUP moves per | |
| question in a chat conversation. It replies with a move copied from the list. With more moves the piece plays out as a | |
| knockout: groups of up to GROUP, then the group winners. | |
| The reply is read, not generated, so there is nothing to parse and no invalid answer. Each question is asked twice, | |
| "Best move?" and "Worst move?", in the same conversation, and a move's score is how much more likely the model is to | |
| reply with it to the first than to the second (average log-probability per token). That mirrors Base mode's | |
| "good move" minus "bad move": what makes a reply likely no matter what was asked (a familiar phrase, a move that | |
| repeats words of the conversation, its place in the list) cancels out, and what is left is whether the model picks | |
| moves according to the question. Two pilots on the pool's models led here: letter answers (A/B/C/D) were close to | |
| chance level, and scoring only "Best move?" replies mostly measured which phrases a model likes (a model picked "good" | |
| moves just as often with no rules and no example in the conversation). | |
| The conversation is multi-turn: two fixed example exchanges (the same for every model), then the model's own previous | |
| question and reply, then the current question, which starts with what the last move did. | |
| """ | |
| from __future__ import annotations | |
| import copy | |
| import random | |
| import torch | |
| from players import ModelRejected | |
| SYSTEM = ( | |
| "You play Tetris. Each turn shows the next piece and a few moves.\n" | |
| "Good moves clear lines, make no holes and keep the stack low and flat.\n" | |
| "Reply with the move asked for, copied exactly." | |
| ) | |
| GROUP = 8 # moves per question | |
| BEST, WORST = "Best move?", "Worst move?" | |
| # First exchanges of every conversation, identical for every model: one "best" and one "worst" question with obvious | |
| # answers, so the model sees both kinds of question and the format of a reply before its first real move. | |
| EXAMPLES = [ | |
| ("Piece: O\n- no lines, 2 new holes, high stack, bumpy\n- clears 1 line, no new holes, low stack, flat\n" + BEST, | |
| "clears 1 line, no new holes, low stack, flat"), | |
| ("Piece: S\n- no lines, no new holes, low stack, flat\n- no lines, 1 new hole, high stack, uneven\n" + WORST, | |
| "no lines, 1 new hole, high stack, uneven"), | |
| ] | |
| # Repos whose tokenizer has no chat template but whose model card documents the prompt format they were tuned on. | |
| _ALPACA = ( | |
| "{{ bos_token }}Below is an instruction that describes a task. Write a response that appropriately completes the " | |
| "request.\n\n{% set ns = namespace(sys='') %}{% for m in messages %}{% if m['role'] == 'system' %}" | |
| "{% set ns.sys = m['content'] + '\n\n' %}{% elif m['role'] == 'user' %}" | |
| "{{ '### Instruction:\n' + ns.sys + m['content'] + '\n\n### Response:\n' }}{% set ns.sys = '' %}" | |
| "{% else %}{{ m['content'] + eos_token + '\n\n' }}{% endif %}{% endfor %}" | |
| ) | |
| CHAT_TEMPLATES = { | |
| "supralabs/supra-50m-instruct": _ALPACA, # Alpaca format (see the repo's sft.py) | |
| } | |
| class NotChatModel(ModelRejected): | |
| """The model can't play Instruct (no chat template, or the conversation does not fit its context).""" | |
| # ---------------------------------------------------------------------------- | |
| # Moves and questions | |
| # ---------------------------------------------------------------------------- | |
| _LINES = {0: "no lines", 1: "clears 1 line", 2: "clears 2 lines", 3: "clears 3 lines", 4: "clears 4 lines"} | |
| def short(c) -> str: | |
| """Outcome of a placement in four facts, always in the same order: lines, holes, height, surface.""" | |
| if c.new_holes < 0: | |
| holes = "removes holes" | |
| elif c.new_holes == 0: | |
| holes = "no new holes" | |
| elif c.new_holes == 1: | |
| holes = "1 new hole" | |
| elif c.new_holes == 2: | |
| holes = "2 new holes" | |
| else: | |
| holes = "many new holes" | |
| mh = c.max_height | |
| height = ("very low stack" if mh <= 4 else "low stack" if mh <= 8 else "high stack" if mh <= 12 | |
| else "very high stack" if mh <= 16 else "stack at the top") | |
| surface = "flat" if c.bump <= 4 else "uneven" if c.bump <= 10 else "bumpy" | |
| return f"{_LINES[c.lines]}, {holes}, {height}, {surface}" | |
| def question(piece: str, moves, last: str | None = None, ask: str = BEST) -> str: | |
| lines = [f"Last move: {last}."] if last else [] | |
| lines.append(f"Piece: {piece}") | |
| lines += [f"- {m}" for m in moves] | |
| lines.append(ask) | |
| return "\n".join(lines) | |
| def knockout(game, seed: int, cands, decide): | |
| """Pick a placement for the current piece. `decide(history, piece, last, moves)` returns one score per move. | |
| Every random draw is seeded by (seed, piece index), so all players see the moves of a piece in the same order. | |
| Returns (candidate, score, reply, number of moves); the conversation state lives in game.memory.""" | |
| memo = game.memory.setdefault("instruct", {"prev": None, "last": None}) | |
| by_text = {} | |
| for c in cands: | |
| by_text.setdefault(short(c), []).append(c) | |
| moves = list(by_text) | |
| random.Random(f"{seed}:{game.pieces}:moves").shuffle(moves) | |
| history = EXAMPLES + ([memo["prev"]] if memo["prev"] else []) | |
| final = None # (question, reply, score) of the last question asked for this piece | |
| rnd = 0 | |
| while len(moves) > 1: | |
| n_groups = -(-len(moves) // GROUP) | |
| groups = [moves[i::n_groups] for i in range(n_groups)] # balanced: sizes differ by one at most | |
| winners = [] | |
| for gi, group in enumerate(groups): | |
| if len(group) == 1: | |
| winners.append(group[0]) | |
| continue | |
| scores = decide(history, game.current_piece, memo["last"], group) | |
| top = max(scores) | |
| ties = [i for i, s in enumerate(scores) if s >= top - 1e-9] | |
| pick = ties[random.Random(f"{seed}:{game.pieces}:{rnd}:{gi}").randrange(len(ties))] | |
| winners.append(group[pick]) | |
| final = (question(game.current_piece, group, memo["last"]), group[pick], scores[pick]) | |
| moves = winners | |
| rnd += 1 | |
| text = moves[0] | |
| same = by_text[text] | |
| cand = same[random.Random(f"{seed}:{game.pieces}").randrange(len(same))] | |
| if final: | |
| memo["prev"] = final[:2] | |
| memo["last"] = text | |
| return cand, (final[2] if final else None), (final[1] if final else ""), len(by_text) | |
| # ---------------------------------------------------------------------------- | |
| # Talking to the model | |
| # ---------------------------------------------------------------------------- | |
| def template_of(player): | |
| """The chat template a player uses (its tokenizer's, or a documented format), or None.""" | |
| return CHAT_TEMPLATES.get(player.model_id.lower()) or getattr(player.tok, "chat_template", None) | |
| def is_chat(player) -> bool: | |
| return getattr(player, "tok", None) is not None and bool(template_of(player)) | |
| def messages(history, user: str): | |
| """history: [(question, reply), ...] before the current question.""" | |
| out = [{"role": "system", "content": SYSTEM}] | |
| for u, a in history: | |
| out += [{"role": "user", "content": u}, {"role": "assistant", "content": a}] | |
| out.append({"role": "user", "content": user}) | |
| return out | |
| def render(player, msgs) -> str: | |
| """The conversation as the model's prompt, ending where its reply starts. Templates without a system role get the | |
| system text at the start of the first user message.""" | |
| tpl = template_of(player) | |
| try: | |
| text = player.tok.apply_chat_template(msgs, chat_template=tpl, tokenize=False, add_generation_prompt=True) | |
| if SYSTEM.splitlines()[0] in text: | |
| return text | |
| except Exception: | |
| pass | |
| merged = [dict(m) for m in msgs if m["role"] != "system"] | |
| merged[0]["content"] = SYSTEM + "\n\n" + merged[0]["content"] | |
| return player.tok.apply_chat_template(merged, chat_template=tpl, tokenize=False, add_generation_prompt=True) | |
| def _encode(tok, text): | |
| return list(tok(text, add_special_tokens=False)["input_ids"]) | |
| def _reply_ids(tok, text, ids, reply): | |
| """Token ids of `reply` right after the prompt `text` (tokenized as ids).""" | |
| full = _encode(tok, text + reply) | |
| return full[len(ids):] if len(full) > len(ids) and full[:len(ids)] == ids else _encode(tok, reply) | |
| def _full(player, common, rows): | |
| """rows: [(suffix, reply)] after the shared prefix `common`. Average log-probability per token of each reply.""" | |
| raw = player._logprob_batch([(common + s + r, len(common) + len(s)) for s, r in rows]) | |
| return [x / len(r) for x, (_, r) in zip(raw, rows)] | |
| def _opening(player): | |
| """Token ids and KV cache of the start that every conversation shares (system prompt and examples), or False.""" | |
| tok = player.tok | |
| a = _encode(tok, render(player, messages(EXAMPLES, question("T", ["x"])))) | |
| b = _encode(tok, render(player, messages(EXAMPLES + [(question("Z", ["y"], "w"), "y")], question("I", ["q"], "v")))) | |
| n = 0 | |
| while n < min(len(a), len(b)) and a[n] == b[n]: | |
| n += 1 | |
| if n < 8: | |
| return False | |
| past = player.model(input_ids=torch.tensor([a[:n]]), attention_mask=torch.ones((1, n), dtype=torch.long), | |
| use_cache=True).past_key_values | |
| return (a[:n], past) if past is not None else False | |
| def _start(player, common): | |
| """KV cache after `common`: the shared opening comes from a copy of its cache, computed once per model.""" | |
| model = player.model | |
| if "_instruct_opening" not in player.__dict__: | |
| player._instruct_opening = _opening(player) | |
| opening = player._instruct_opening | |
| if opening and len(common) > len(opening[0]) and common[: len(opening[0])] == opening[0]: | |
| rest = torch.tensor([common[len(opening[0]):]]) | |
| mask = torch.ones((1, len(common)), dtype=torch.long) | |
| return model(input_ids=rest, attention_mask=mask, past_key_values=copy.deepcopy(opening[1]), use_cache=True).past_key_values | |
| inp = torch.tensor([common]) | |
| return model(input_ids=inp, attention_mask=torch.ones_like(inp), use_cache=True).past_key_values | |
| def _cached(player, common, rows): | |
| """Same as _full, computing the shared prefix once and every suffix + reply together on its KV cache.""" | |
| model = player.model | |
| past = _start(player, common) | |
| past.batch_repeat_interleave(len(rows)) # AttributeError on caches that can't be shared: caller falls back | |
| width = max(len(s) + len(r) for s, r in rows) | |
| ids = torch.full((len(rows), width), player.pad_id, dtype=torch.long) | |
| mask = torch.zeros((len(rows), len(common) + width), dtype=torch.long) | |
| mask[:, : len(common)] = 1 | |
| for i, (s, r) in enumerate(rows): | |
| ids[i, : len(s) + len(r)] = torch.tensor(s + r) | |
| mask[i, len(common): len(common) + len(s) + len(r)] = 1 | |
| lp = torch.log_softmax(model(input_ids=ids, attention_mask=mask, past_key_values=past, use_cache=False).logits.float(), dim=-1) | |
| out = [] | |
| for i, (s, r) in enumerate(rows): | |
| out.append(sum(lp[i, len(s) + j - 1, r[j]].item() for j in range(len(r))) / len(r)) | |
| return out | |
| def _copied(player, common, rows): | |
| """Same as _full for caches that can't be batched (e.g. hybrid convolution layers): the shared prefix is computed | |
| once, then every suffix + reply runs on its own copy of the cache.""" | |
| model = player.model | |
| past = _start(player, common) | |
| if past is None: | |
| raise ValueError("the model returns no cache") | |
| out = [] | |
| for s, r in rows: | |
| ids = torch.tensor([s + r]) | |
| mask = torch.ones((1, len(common) + len(s) + len(r)), dtype=torch.long) | |
| logits = model(input_ids=ids, attention_mask=mask, past_key_values=copy.deepcopy(past), use_cache=True).logits | |
| lp = torch.log_softmax(logits[0].float(), dim=-1) | |
| out.append(sum(lp[len(s) + j - 1, r[j]].item() for j in range(len(r))) / len(r)) | |
| return out | |
| _WAYS = {"batch": _cached, "copy": _copied, "full": _full} | |
| def contrast_scores(player, best_text: str, worst_text: str, moves) -> list[float]: | |
| """For each move: average log-probability per token of replying with it to the "best" prompt minus the same for the | |
| "worst" prompt. The two prompts differ only at the end, so the shared prefix is computed once.""" | |
| tok = player.tok | |
| b, w = _encode(tok, best_text), _encode(tok, worst_text) | |
| n = 0 | |
| while n < min(len(b), len(w)) - 1 and b[n] == w[n]: | |
| n += 1 | |
| common = b[:n] | |
| rows = [(b[n:], _reply_ids(tok, best_text, b, m)) for m in moves] + [(w[n:], _reply_ids(tok, worst_text, w, m)) for m in moves] | |
| longest = max(len(b), len(w)) + max(len(r) for _, r in rows) | |
| if longest > player.max_len: | |
| raise NotChatModel(f"{player.model_id}: the conversation ({longest} tokens) exceeds its context ({player.max_len}).") | |
| way = player.__dict__.get("_instruct_way") | |
| if way is None: # first question: use the fastest way that gives the same scores as the full computation | |
| scores = _full(player, common, rows) | |
| player._instruct_way = "full" | |
| for name in ("batch", "copy"): | |
| try: | |
| if max(abs(a - c) for a, c in zip(_WAYS[name](player, common, rows), scores)) < 1e-3: | |
| player._instruct_way = name | |
| break | |
| except Exception: | |
| continue | |
| else: | |
| scores = _WAYS[way](player, common, rows) | |
| k = len(moves) | |
| return [scores[i] - scores[k + i] for i in range(k)] | |
| def _prompts(player, history, piece, last, moves): | |
| """The "best" and "worst" prompts; drops the previous exchange (never the examples) if they would not fit.""" | |
| history = list(history) | |
| while True: | |
| best = render(player, messages(history, question(piece, moves, last, BEST))) | |
| worst = render(player, messages(history, question(piece, moves, last, WORST))) | |
| if len(_encode(player.tok, best)) + 40 <= player.max_len or len(history) <= len(EXAMPLES): | |
| return best, worst | |
| history.pop() | |
| def ask(player, history, piece, last, moves) -> list[float]: | |
| with player._lock: | |
| return contrast_scores(player, *_prompts(player, history, piece, last, moves), list(moves)) | |
| def check(player) -> None: | |
| """Raise NotChatModel unless the player can hold an Instruct conversation (once per loaded model).""" | |
| if player.__dict__.get("_instruct_ok"): | |
| return | |
| if not is_chat(player): | |
| raise NotChatModel(f"`{player.model_id}` has no chat template: Instruct mode is for chat (instruct) models.") | |
| longest = [f"clears {i % 4 + 1} lines, many new holes, very high stack, bumpy" for i in range(GROUP)] | |
| prev = (question("T", longest, "no lines, removes holes, stack at the top, uneven"), longest[0]) | |
| try: | |
| ask(player, EXAMPLES + [prev], "T", "no lines, removes holes, stack at the top, uneven", longest) | |
| except NotChatModel: | |
| raise | |
| except Exception as e: | |
| raise NotChatModel(f"`{player.model_id}`: the chat prompt failed ({type(e).__name__}: {str(e)[:200]}).") | |
| player._instruct_ok = True | |
| def choose(game, player, seed: int, cands): | |
| """One Instruct move for a chat model: plays it on `game`.""" | |
| cand, score, reply, n = knockout(game, seed, cands, lambda h, piece, last, moves: ask(player, h, piece, last, moves)) | |
| game.apply(cand, score, n, reply) | |
| return True | |
| def transcript() -> str: | |
| """The start of an Instruct conversation as plain text (How it works tab).""" | |
| q = question("T", ["no lines, no new holes, low stack, flat", "no lines, 1 new hole, low stack, bumpy", | |
| "clears 2 lines, no new holes, low stack, flat", "no lines, no new holes, high stack, flat"], | |
| last="clears 1 line, no new holes, low stack, flat") | |
| parts = [("system", SYSTEM)] | |
| for u, a in EXAMPLES: | |
| parts += [("user", u), ("assistant", a)] | |
| parts.append(("user", q)) | |
| return "\n\n".join(f"[{role}]\n{text}" for role, text in parts) | |