Add more handlers

app.py

@@ -1,11 +1,12 @@
 import logging
+
 from src import interface
 
 logging.basicConfig(level="DEBUG")
 
 
 def main():
-    demo = interface.get_demo(wip=True)
+    demo = interface.get_demo(wip=False)
     demo.launch(debug=True)
 
 

src/handler.py

@@ -27,15 +27,26 @@ def get_model_predictions(input_text: str) -> torch.Tensor:
     return top_3
 
 
-def guess_is_correct(text: str, next_token: int) -> bool:
+def guess_is_correct(text: str) -> bool:
     """
     We check if the predicted token or a corresponding one with a leading whitespace
     matches that of the next token
     """
-    logger.debug("Next token: '{}'".format(tokenizer.convert_ids_to_tokens([next_token])))
+    current_target = all_tokens[STATE.current_word_index]
+    logger.debug("Next token: '{}'".format(tokenizer.convert_ids_to_tokens([current_target])))
     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)
+    return current_target in (predicted_token_start, predicted_token_whitespace)
+
+
+def lm_is_correct() -> bool:
+    # NOTE: out of range if remaining attempts is 0
+    if STATE.remaining_attempts > 1:
+        return False
+
+    current_guess = STATE.lm_guesses[MAX_ATTEMPTS - STATE.remaining_attempts]
+    current_target = all_tokens[STATE.current_word_index]
+    return current_guess == current_target
 
 
 def get_start_and_whitespace_tokens(word: str) -> tuple[int, int]:
@@ -44,73 +55,135 @@ def get_start_and_whitespace_tokens(word: str) -> tuple[int, int]:
     return predicted_token_start, predicted_token_whitespace
 
 
+def get_current_text():
+    return tokenizer.decode(all_tokens[: STATE.current_word_index])
+
+
+def handle_player_win():
+    # TODO: point system
+    points = 1
+    STATE.player_points += points
+    STATE.button_label = "Next word"
+    return STATE.get_tuple(
+        get_current_text(),
+        bottom_html=f"Player gets {points} point!",
+    )
+
+
+def handle_lm_win():
+    points = 1
+    STATE.lm_points += points
+    STATE.button_label = "Next word"
+    return STATE.get_tuple(
+        get_current_text(),
+        bottom_html=f"GPT2 gets {points} point!",
+    )
+
+
+def handle_out_of_attempts():
+    STATE.button_label = "Next word"
+    return STATE.get_tuple(
+        get_current_text(),
+        bottom_html="Out of attempts. No one gets points!",
+    )
+
+
+def handle_tie():
+    STATE.button_label = "Next word"
+    return STATE.get_tuple(
+        get_current_text(),
+        bottom_html="TIE! No one gets points!",
+    )
+
+
+def handle_next_attempt():
+    STATE.remaining_attempts -= 1
+    return STATE.get_tuple(
+        get_current_text(), bottom_html=f"That was not it... {STATE.remaining_attempts} attempts left"
+    )
+
+
+def handle_no_input():
+    return STATE.get_tuple(
+        get_current_text(),
+        bottom_html="Please write something",
+    )
+
+
+def handle_next_word():
+    STATE.next_word()
+    STATE.lm_guesses = get_model_predictions(get_current_text())
+    return STATE.get_tuple()
+
+
 def handle_guess(
     text: str,
-    remaining_attempts: int,
     *args,
     **kwargs,
 ) -> str:
     """
     * Retreives model predictions and compares the top 3 predicted tokens
     """
-    logger.debug(
-        f"Params:\ntext = {text}\n"
-        f"remaining_attempts = {remaining_attempts}\n"
-        f"args = {args}\n"
-        f"kwargs = {kwargs}\n"
-    )
+    logger.debug("Params:\n" f"text = {text}\n" f"args = {args}\n" f"kwargs = {kwargs}\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 = ""
+    if STATE.button_label == "Next word":
+        return handle_next_word()
 
     if not text:
-        logger.debug("Returning early")
-        return (
-            current_text, 
-            STATE.player_points,
-            STATE.lm_points,
-            STATE.player_guess_str,
-            STATE.get_lm_guess_display(remaining_attempts),
-            remaining_attempts,
-            "",
-            "Guess!"
-        )
-
-    if remaining_attempts == 0:
-        STATE.next_word()
-        current_tokens = all_tokens[: STATE.current_word_index]
-        remaining_attempts = MAX_ATTEMPTS
-
-    remaining_attempts -= 1
-
-    next_token = all_tokens[STATE.current_word_index]
-
-    if guess_is_correct(text, next_token):
-        # STATE.correct_guess()
-        STATE.player_points += 1
-        remaining_attempts = 0
+        return handle_no_input()
 
+    STATE.player_guesses.append(text)
+
+    player_correct = guess_is_correct(text)
+    lm_correct = lm_is_correct()
+
+    if player_correct and lm_correct:
+        return handle_tie()
+    elif player_correct and not lm_correct:
+        return handle_player_win()
+    elif lm_correct and not player_correct:
+        return handle_lm_win()
+    elif STATE.remaining_attempts == 0:
+        return handle_out_of_attempts()
     else:
-        STATE.player_guesses.append(text)
-
-    # FIXME: unoptimized, computing all three every time
-    current_text = tokenizer.decode(current_tokens)
-    STATE.lm_guesses = get_model_predictions(current_text)[: MAX_ATTEMPTS - remaining_attempts]
-    logger.debug(f"lm_guesses: {tokenizer.decode(STATE.lm_guesses)}")
-    logger.debug(f"Pre-return STATE:\n{STATE}")
-
-    # BUG: if you enter the word guess field when it says next
-    #      word, it will guess it as the next
-    return (
-        current_text,
-        STATE.player_points,
-        STATE.lm_points,
-        STATE.player_guess_str,
-        STATE.get_lm_guess_display(remaining_attempts),
-        remaining_attempts,
-        "",
-        "Guess!" if remaining_attempts else "Next word",
-    )
+        return handle_next_attempt()
+
+
+STATE.lm_guesses = get_model_predictions(get_current_text())
+
+
+#     # STATE.correct_guess()
+#     # remaining_attempts = 0
+# # elif lm_guess_is_correct():
+# #     pass
+# else:
+#     return handle_incorrect_guess()
+# # elif remaining_attempts == 0:
+# #     return handle_out_of_attempts()
+
+#     remaining_attempts -= 1
+#     STATE.player_guesses.append(text)
+
+# if remaining_attempts == 0:
+#     STATE.next_word()
+#     current_tokens = all_tokens[: STATE.current_word_index]
+#     remaining_attempts = MAX_ATTEMPTS
+
+# # FIXME: unoptimized, computing all three every time
+# current_text = tokenizer.decode(current_tokens)
+# logger.debug(f"lm_guesses: {tokenizer.decode(STATE.lm_guesses)}")
+# logger.debug(f"Pre-return STATE:\n{STATE}")
+
+# # BUG: if you enter the word guess field when it says next
+# #      word, it will guess it as the next
+# return (
+#     current_text,
+#     STATE.player_points,
+#     STATE.lm_points,
+#     STATE.player_guess_str,
+#     STATE.get_lm_guess_display(remaining_attempts),
+#     remaining_attempts,
+#     "",
+#     "Guess!" if remaining_attempts else "Next word",
+# )

src/interface.py

@@ -7,7 +7,11 @@ from src.state import STATE
 from src.state import tokenizer
 
 
-def build_demo():    
+def build_demo():
+    """
+    Builds and returns the gradio app interface
+    """
+
     with gr.Blocks() as demo:
         with gr.Row():
             gr.Markdown("<h1><center>Can you beat a language model?</center></h1>")
@@ -35,6 +39,7 @@ def build_demo():
                     value=MAX_ATTEMPTS,
                     label="Remaining attempts",
                     precision=0,
+                    interactive=False,
                 )
                 current_guesses = gr.Textbox(label="Your guesses")
             with gr.Column():
@@ -45,11 +50,13 @@ def build_demo():
                 guess_field = gr.Textbox(label="")
                 guess_button = gr.Button(value="Guess!")
 
+        with gr.Row():
+            bottom_html = gr.HTML()
+
         guess_button.click(
             handle_guess,
             inputs=[
                 guess_field,
-                remaining_attempts,
             ],
             outputs=[
                 prompt_text,
@@ -60,6 +67,7 @@ def build_demo():
                 remaining_attempts,
                 guess_field,
                 guess_button,
+                bottom_html,
             ],
         )
 

src/state.py

@@ -1,3 +1,4 @@
+import logging
 from dataclasses import dataclass
 
 from transformers import AutoModelForCausalLM
@@ -5,31 +6,64 @@ from transformers import AutoTokenizer
 
 from src.constants import MAX_ATTEMPTS
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass
 class ProgramState:
     current_word_index: int
-    player_guesses: list
+    player_guesses: list[str]
     player_points: int
-    lm_guesses: list
+    lm_guesses: list[int]
     lm_points: int
+    remaining_attempts: int
+    button_label: str
+
+    @property
+    def player_guess_str(self):
+        return "\n".join(self.player_guesses)
+
+    @property
+    def lm_guess_str(self):
+        strings = list(map(tokenizer.decode, self.lm_guesses))
+        logger.debug(strings)
+
+        n_censored = self.remaining_attempts
+        for i in range(1, n_censored + 1):
+            strings[-i] = "****"
 
-    def correct_guess(self):
-        # FIXME: not 1 for every point
-        self.player_points += 1
-        self.next_word()
+        logger.debug(strings)
+        return "\n".join(strings)
 
     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)
+        self.lm_guesses = []  # TODO: make guesses?
+        self.button_label = "Guess!"
 
-    def get_lm_guess_display(self, remaining_attempts: int) -> str:
-        return "\n".join(map(tokenizer.decode, self.lm_guesses[: MAX_ATTEMPTS - remaining_attempts]))
+    def get_tuple(
+        self,
+        prompt_text=None,
+        player_points=None,
+        lm_points=None,
+        player_guess_str=None,
+        lm_guess_str=None,
+        remaining_attempts=None,
+        text_field=None,
+        button_label=None,
+        bottom_html=None,
+    ) -> tuple:
+        return (
+            prompt_text or "",  # FIXME
+            player_points or self.player_points,
+            lm_points or self.lm_points,
+            player_guess_str or self.player_guess_str,
+            lm_guess_str or self.lm_guess_str,
+            remaining_attempts or self.remaining_attempts,
+            text_field or "",  # FIXME
+            button_label or self.button_label,
+            bottom_html or "",  # FIXME
+        )
 
 
 STATE = ProgramState(
@@ -38,6 +72,8 @@ STATE = ProgramState(
     lm_guesses=[],
     player_points=0,
     lm_points=0,
+    remaining_attempts=MAX_ATTEMPTS,
+    button_label="Guess!",
 )
 
 tokenizer = AutoTokenizer.from_pretrained("gpt2")

src/utils.py

@@ -0,0 +1,20 @@
+from collections import namedtuple
+
+
+HandlerInput = namedtuple(
+    typename="HandlerInput", 
+    field_names=[
+        "text",
+        "remaining_attempts",
+        "button_label",
+    ],
+)
+
+HandlerOutput = namedtuple(
+    typename="HandlerOutput", 
+    field_names=[
+        "text",
+        "remaining_attempts",
+        "button_label",
+    ],
+)