LMPlayer.complete: the model's raw greedy completion for the chosen move

players.py

@@ -16,6 +16,7 @@ MIN_PARAMS = int(os.environ.get("MIN_PARAMS", 50_000))
 ALLOW_REMOTE_CODE = os.environ.get("ALLOW_REMOTE_CODE", "1") == "1"
 MAX_CACHED_MODELS = int(os.environ.get("MAX_CACHED_MODELS", 6))
 BATCH_SIZE = 16
+GEN_TOKENS = 8  # length of the raw completion shown under each board
 
 PROMPTS = {
     # The rules of the game are stated in the prompt: tests reading comprehension.
@@ -113,6 +114,8 @@ class LMPlayer:
         self.n_params = n_params
         self.custom_code = custom_code
         self._cache: dict[tuple[str, str], float] = {}
+        self._gen_cache: dict[tuple[str, str], str] = {}
+        self._use_generate = None  # decided on the first completion
         self._lock = threading.Lock()
         self._fwd_kwargs = None  # discovered on first forward
         self.max_len = self._max_len()
@@ -217,6 +220,43 @@ class LMPlayer:
         table = self.score_descriptions(protocol, [c.description for c in cands])
         return [table[c.description] for c in cands]
 
+    @torch.inference_mode()
+    def complete(self, protocol, desc):
+        """The model's own greedy text after the prompt of the chosen move (display only, not used to decide)."""
+        key = (protocol, desc)
+        with self._lock:
+            if key in self._gen_cache:
+                return self._gen_cache[key]
+            ids = self.prefix + self._ids(PROMPTS[protocol].format(desc=desc), False)
+            n_new = min(GEN_TOKENS, self.max_len - len(ids))
+            out = self._generate(ids, n_new) if n_new > 0 else []
+            text = self.tok.decode(out, skip_special_tokens=True).split("\n")[0]
+            self._gen_cache[key] = text
+            return text
+
+    def _generate(self, ids, n_new):
+        """Greedy decoding: `generate()` (with KV cache) when the model supports it, else a plain loop."""
+        inp = torch.tensor([ids], dtype=torch.long)
+        if self._use_generate is not False:
+            try:
+                gen = self.model.generate(inp, attention_mask=torch.ones_like(inp), max_new_tokens=n_new, do_sample=False,
+                                          pad_token_id=self.pad_id, eos_token_id=self.tok.eos_token_id)
+                self._use_generate = True
+                return gen[0, len(ids):].tolist()
+            except Exception:
+                self._use_generate = False
+        eos = self.tok.eos_token_id
+        out = []
+        for _ in range(n_new):
+            inp = torch.tensor([ids + out], dtype=torch.long)
+            nxt = int(self._forward(inp, torch.ones_like(inp))[0, -1].argmax())
+            if nxt == eos:
+                break
+            out.append(nxt)
+            if "\n" in self.tok.decode(out, skip_special_tokens=True):
+                break
+        return out
+
 
 # ----------------------------------------------------------------------------
 # Validation + loading