WIP: Add attempt count

app.py

@@ -1,75 +1,11 @@
-import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM
-from transformers import AutoTokenizer
+import logging
+from src import interface
 
-from state import ProgramState
-from text import get_text
-
-STATE = ProgramState(
-    current_token=20,
-)
-
-
-tokenizer = AutoTokenizer.from_pretrained("gpt2")
-model = AutoModelForCausalLM.from_pretrained("gpt2")
-model.eval()
-
-all_tokens = tokenizer.encode(get_text())
-
-
-def handle_guess(text: str) -> str:
-    """
-    Retreives 
-    """
-    STATE.current_token += 1
-    decoded_str = tokenizer.decode(all_tokens[:STATE.current_token])
-    return decoded_str, ""
-
-
-def get_model_predictions(input_text: str) -> torch.Tensor:
-    """
-    Returns the indices as a torch tensor of the top 3 predicted tokens.
-    """
-    inputs = tokenizer(input_text, return_tensors="pt")
-
-    with torch.no_grad():
-        logits = model(**inputs).logits
-
-    last_token = logits[0, -1]
-    top_3 = torch.topk(last_token, 3)
-
-    return top_3
-
-def build_demo():
-    with gr.Blocks() as demo:
-        gr.Markdown("<h1><center>Can you beat a language model?</center></h1>")
-
-        with gr.Row():
-            prompt_text = gr.Markdown()
-        with gr.Row():
-            with gr.Column():
-                guess = gr.Textbox(label="Guess!")
-                guess_btn = gr.Button(value="Guess!")
-            with gr.Column():
-                lm_guess = gr.Textbox(label="LM guess")
-
-        guess_btn.click(handle_guess, inputs=guess, outputs=[prompt_text, lm_guess], api_name="get_next_word")
-    return demo
-
-
-def wip_sign():
-    with gr.Blocks() as demo:
-        gr.Markdown("<h1><center>Can you beat a language model?</center></h1>")
-        with gr.Row():
-            gr.Markdown("<h1><center>⛔👷‍♂️ Work in progress, come back later </center></h1>")
-
-    return demo
+logging.basicConfig(level="DEBUG")
 
 
 def main():
-    demo = wip_sign()
-    # demo = build_demo()
+    demo = interface.get_demo(wip=True)
     demo.launch(debug=True)
 
 

src/handler.py

@@ -0,0 +1,77 @@
+import torch
+import logging
+
+from src.state import STATE
+from src.state import tokenizer
+from src.state import model
+from src.text import get_text
+
+logger = logging.getLogger(__name__)
+
+all_tokens = tokenizer.encode(get_text())
+
+
+def get_model_predictions(input_text: str) -> torch.Tensor:
+    """
+    Returns the indices as a torch tensor of the top 3 predicted tokens.
+    """
+    inputs = tokenizer(input_text, return_tensors="pt")
+
+    with torch.no_grad():
+        logits = model(**inputs).logits
+
+    last_token = logits[0, -1]
+    top_3 = torch.topk(last_token, 3).indices.tolist()
+    return top_3
+
+
+def handle_guess(text: str) -> str:
+    """
+    * 
+    * Retreives model predictions and compares the top 3 predicted tokens 
+    """
+    current_tokens = all_tokens[:STATE.current_word_index]
+    current_text = tokenizer.decode(current_tokens)
+    player_guesses = ""
+    lm_guesses = ""
+    remaining_attempts = 3
+
+    if not text:
+        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 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 = []
+    else:
+        remaining_attempts -= 1
+        STATE.player_guesses.append(tokenizer.decode([predicted_token_whitespace]))
+
+    # 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)
+
+    return (
+        current_text,
+        player_guesses,
+        lm_guesses,
+        remaining_attempts
+    )

src/interface.py

@@ -0,0 +1,68 @@
+import gradio as gr
+from src.handler import handle_guess
+
+
+def build_demo():
+    with gr.Blocks() as demo:
+        with gr.Row():
+            gr.Markdown("<h1><center>Can you beat a language model?</center></h1>")
+        with gr.Row():
+            gr.Markdown(
+                "Can you beat language models at their own game?\n"
+                "In this game you're pitted against a language model in the task of, you guessed it, laungage modelling.\n"
+                "Your task is to predict the next word given the previous sequence. You will get 3 attempts to guess.\n"
+                "The one with the fewest guesses for a given word gets a point."
+            )
+        with gr.Row():
+            prompt_text = gr.Textbox(label="Context", interactive=False)
+        with gr.Row():
+            with gr.Column():
+                player_points = gr.Number(label="your points", interactive=False)
+            with gr.Column():
+                lm_points = gr.Number(label="LM points", interactive=False)
+        with gr.Row():
+            with gr.Column():
+                remaining_attempts = gr.Number(label="Remaining attempts")
+                current_guesses = gr.Textbox(label="Your guesses")
+            with gr.Column(): 
+                lm_guesses = gr.Textbox(label="LM guesses")
+
+        with gr.Row():
+            with gr.Column():
+                guess = 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,
+            outputs=[
+                prompt_text,
+                current_guesses,
+                lm_guesses,
+                remaining_attempts,
+            ],
+        )
+
+    return demo
+
+
+def wip_sign():
+    with gr.Blocks() as demo:
+        gr.Markdown("<h1><center>Can you beat a language model?</center></h1>")
+        with gr.Row():
+            gr.Markdown("<h1><center>⛔👷‍♂️ Work in progress, come back later </center></h1>")
+
+    return demo
+
+
+def get_demo(wip=False):
+    if wip:
+        return wip_sign()
+    else:
+        return build_demo()
+
+
+

src/state.py

@@ -0,0 +1,27 @@
+from dataclasses import dataclass
+from transformers import AutoTokenizer
+from transformers import AutoModelForCausalLM
+
+
+from dataclasses import dataclass
+
+@dataclass
+class ProgramState:
+    current_word_index: int
+    player_guesses: list
+    player_points: int
+    lm_guesses: list
+    lm_points: int
+
+
+STATE = ProgramState(
+    current_word_index=20,
+    player_guesses=[],
+    lm_guesses=[],
+    player_points=0,
+    lm_points=0,
+)
+
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+model = AutoModelForCausalLM.from_pretrained("gpt2")
+model.eval()

state.py

@@ -1,10 +0,0 @@
-from dataclasses import dataclass
-
-
-@dataclass
-class ProgramState:
-    current_token: int
-    # full_text: str
-
-    def get_text(self):
-        pass