Add attempt counts

src/constants.py

@@ -0,0 +1 @@
+MAX_ATTEMPTS = 3

src/handler.py

@@ -1,9 +1,11 @@
-import torch
 import logging
 
+import torch
+
+from src.constants import MAX_ATTEMPTS
 from src.state import STATE
-from src.state import tokenizer
 from src.state import model
+from src.state import tokenizer
 from src.text import get_text
 
 logger = logging.getLogger(__name__)
@@ -25,53 +27,65 @@ def get_model_predictions(input_text: str) -> torch.Tensor:
     return top_3
 
 
-def handle_guess(text: str) -> str:
+def guess_is_correct(text: str, next_token: int) -> bool:
     """
-    * 
-    * Retreives model predictions and compares the top 3 predicted tokens 
+    We check if the predicted token or a corresponding one with a leading whitespace
+    matches that of the next token
     """
-    current_tokens = all_tokens[:STATE.current_word_index]
+    logger.debug("Next token: '{}'".format(tokenizer.convert_ids_to_tokens([next_token])))
+    predicted_token_start, predicted_token_whitespace = get_start_and_whitespace_tokens(text)
+    logger.debug(tokenizer.convert_ids_to_tokens([predicted_token_start, predicted_token_whitespace]))
+    return next_token in (predicted_token_start, predicted_token_whitespace)
+
+
+def get_start_and_whitespace_tokens(word: str) -> tuple[int, int]:
+    predicted_token_start = tokenizer.encode(word, add_special_tokens=False)[0]
+    predicted_token_whitespace = tokenizer.encode(". " + word, add_special_tokens=False)[1]
+    return predicted_token_start, predicted_token_whitespace
+
+
+def handle_guess(
+    text: str,
+    remaining_attempts: int,
+) -> str:
+    """
+    *
+    * Retreives model predictions and compares the top 3 predicted tokens
+    """
+    logger.debug(f"Params:\ntext = {text}\nremaining_attempts = {remaining_attempts}\n")
+    logger.debug(f"Initial STATE:\n{STATE}")
+
+    current_tokens = all_tokens[: STATE.current_word_index]
     current_text = tokenizer.decode(current_tokens)
     player_guesses = ""
     lm_guesses = ""
-    remaining_attempts = 3
+    remaining_attempts -= 1
 
     if not text:
-        return (
-            current_text,
-            player_guesses,
-            lm_guesses,
-            remaining_attempts
-        )
+        logger.debug("Returning early")
+        return (current_text, player_guesses, lm_guesses, remaining_attempts)
 
     next_token = all_tokens[STATE.current_word_index]
-    predicted_token_start = tokenizer.encode(text, add_special_tokens=False)[0]
-    predicted_token_whitespace = tokenizer.encode(". " + text, add_special_tokens=False)[1]
-    logger.debug("Next token: '{}'".format(tokenizer.convert_ids_to_tokens([next_token])))
-    logger.debug(tokenizer.convert_ids_to_tokens([predicted_token_start, predicted_token_whitespace]))
 
-    guess_is_correct = next_token in (predicted_token_start, predicted_token_whitespace)
+    if guess_is_correct(text, next_token):
+        STATE.correct_guess()
 
-    if guess_is_correct or remaining_attempts == 0:
-        STATE.current_word_index += 1
-        current_tokens = all_tokens[:STATE.current_word_index]
-        remaining_attempts = 3
-        STATE.player_guesses = []
-        STATE.lm_guesses = []
+    if remaining_attempts == 0:
+        STATE.next_word()
+        current_tokens = all_tokens[: STATE.current_word_index]
+        remaining_attempts = MAX_ATTEMPTS
     else:
-        remaining_attempts -= 1
-        STATE.player_guesses.append(tokenizer.decode([predicted_token_whitespace]))
+        STATE.player_guesses.append(text)
 
     # FIXME: unoptimized, computing all three every time
-    STATE.lm_guesses = get_model_predictions(tokenizer.decode(current_tokens))[:3-remaining_attempts]
-    logger.debug(f"lm_guesses: {tokenizer.decode(lm_guesses)}")
-
-    player_guesses = "\n".join(STATE.player_guesses)
     current_text = tokenizer.decode(current_tokens)
+    STATE.lm_guesses = get_model_predictions(current_text)[: 3 - remaining_attempts]
+    logger.debug(f"lm_guesses: {tokenizer.decode(STATE.lm_guesses)}")
 
+    logger.debug(f"Pre-return STATE:\n{STATE}")
     return (
         current_text,
-        player_guesses,
-        lm_guesses,
-        remaining_attempts
-    )
+        STATE.player_guess_str,
+        STATE.get_lm_guess_display(remaining_attempts),
+        remaining_attempts,
+    )

src/interface.py

@@ -1,5 +1,10 @@
 import gradio as gr
+
+from src.constants import MAX_ATTEMPTS
+from src.handler import all_tokens
 from src.handler import handle_guess
+from src.state import STATE
+from src.state import tokenizer
 
 
 def build_demo():
@@ -14,7 +19,11 @@ def build_demo():
                 "The one with the fewest guesses for a given word gets a point."
             )
         with gr.Row():
-            prompt_text = gr.Textbox(label="Context", interactive=False)
+            prompt_text = gr.Textbox(
+                value=tokenizer.decode(all_tokens[: STATE.current_word_index]),
+                label="Context",
+                interactive=False,
+            )
         with gr.Row():
             with gr.Column():
                 player_points = gr.Number(label="your points", interactive=False)
@@ -22,9 +31,13 @@ def build_demo():
                 lm_points = gr.Number(label="LM points", interactive=False)
         with gr.Row():
             with gr.Column():
-                remaining_attempts = gr.Number(label="Remaining attempts")
+                remaining_attempts = gr.Number(
+                    value=MAX_ATTEMPTS,
+                    label="Remaining attempts",
+                    precision=0,
+                )
                 current_guesses = gr.Textbox(label="Your guesses")
-            with gr.Column(): 
+            with gr.Column():
                 lm_guesses = gr.Textbox(label="LM guesses")
 
         with gr.Row():
@@ -37,7 +50,10 @@ def build_demo():
 
         guess_button.click(
             handle_guess,
-            inputs=guess,
+            inputs=[
+                guess,
+                remaining_attempts,
+            ],
             outputs=[
                 prompt_text,
                 current_guesses,
@@ -63,6 +79,3 @@ def get_demo(wip=False):
         return wip_sign()
     else:
         return build_demo()
-
-
-

src/state.py

@@ -1,9 +1,10 @@
 from dataclasses import dataclass
-from transformers import AutoTokenizer
+
 from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
 
+from src.constants import MAX_ATTEMPTS
 
-from dataclasses import dataclass
 
 @dataclass
 class ProgramState:
@@ -13,6 +14,23 @@ class ProgramState:
     lm_guesses: list
     lm_points: int
 
+    def correct_guess(self):
+        # FIXME: not 1 for every point
+        self.player_points += 1
+        self.next_word()
+
+    def next_word(self):
+        self.current_word_index += 1
+        self.player_guesses = []
+        self.lm_guesses = []
+
+    @property
+    def player_guess_str(self):
+        return "\n".join(self.player_guesses)
+
+    def get_lm_guess_display(self, remaining_attempts: int) -> str:
+        return "\n".join(map(tokenizer.decode, self.lm_guesses[: MAX_ATTEMPTS - remaining_attempts]))
+
 
 STATE = ProgramState(
     current_word_index=20,
@@ -24,4 +42,4 @@ STATE = ProgramState(
 
 tokenizer = AutoTokenizer.from_pretrained("gpt2")
 model = AutoModelForCausalLM.from_pretrained("gpt2")
-model.eval()
+model.eval()

src/text.py

@@ -13,4 +13,4 @@ Nine months later, he was named the 2009 Nobel Peace Prize laureate, a decision
 
 
 def get_text():
-    return target_text
+    return target_text