Update label to "next word"

src/handler.py

@@ -47,45 +47,70 @@ def get_start_and_whitespace_tokens(word: str) -> tuple[int, int]:
 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}\nremaining_attempts = {remaining_attempts}\n")
+    logger.debug(
+        f"Params:\ntext = {text}\n"
+        f"remaining_attempts = {remaining_attempts}\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_tokens = all_tokens[:STATE.current_word_index]
     current_text = tokenizer.decode(current_tokens)
     player_guesses = ""
     lm_guesses = ""
-    remaining_attempts -= 1
 
     if not text:
         logger.debug("Returning early")
-        return (current_text, player_guesses, lm_guesses, remaining_attempts)
-
-    next_token = all_tokens[STATE.current_word_index]
-
-    if guess_is_correct(text, next_token):
-        STATE.correct_guess()
+        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
+
     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)[: 3 - remaining_attempts]
+    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",
     )

src/interface.py

@@ -7,7 +7,7 @@ from src.state import STATE
 from src.state import tokenizer
 
 
-def build_demo():
+def build_demo():    
     with gr.Blocks() as demo:
         with gr.Row():
             gr.Markdown("<h1><center>Can you beat a language model?</center></h1>")
@@ -42,23 +42,24 @@ def build_demo():
 
         with gr.Row():
             with gr.Column():
-                guess = gr.Textbox(label="")
+                guess_field = gr.Textbox(label="")
                 guess_button = gr.Button(value="Guess!")
 
-        with gr.Row():
-            next_word = gr.Button(value="Next word")
-
         guess_button.click(
             handle_guess,
             inputs=[
-                guess,
+                guess_field,
                 remaining_attempts,
             ],
             outputs=[
                 prompt_text,
+                player_points,
+                lm_points,
                 current_guesses,
                 lm_guesses,
                 remaining_attempts,
+                guess_field,
+                guess_button,
             ],
         )