:construction: Make the repo work somewhat

src/handlers.py

@@ -4,6 +4,7 @@ from transformers import PreTrainedTokenizer
 
 from src.params import ReducerParams
 from src.shared import all_tokens
+from src.utils import get_current_prompt_text
 
 logger = logging.getLogger(__name__)
 
@@ -56,4 +57,5 @@ def handle_next_word(params: ReducerParams) -> ReducerParams:
     params.word_number += 1
     params.button_label = "Guess!"
     params.bottom_html = ""
+    params.prompt_text = get_current_prompt_text(params.word_number)
     return params

src/interface.py

@@ -23,7 +23,7 @@ def build_demo():
             )
         with gr.Row():
             prompt_text = gr.Textbox(
-                value=shared.all_tokens[:STARTING_INDEX],
+                value=shared.tokenizer.decode(shared.all_tokens[:STARTING_INDEX]),
                 label="Context",
                 interactive=False,
             )
@@ -52,7 +52,7 @@ def build_demo():
         with gr.Row():
             bottom_html = gr.HTML()
         with gr.Row():
-            word_number = gr.Number(label="Word no.", interactive=False)
+            word_number = gr.Number(label="Word no.", interactive=False, precision=0)
 
         guess_button.click(
             reducer.handle_guess,

src/params.py

@@ -23,3 +23,6 @@ class ReducerParams:
 
     def __iter__(self) -> Iterator:
         return map(partial(getattr, self), self.__dataclass_fields__)
+
+    def __getitem__(self, index) -> str | int:
+        return list(self)[index]

src/predictions.py

@@ -14,35 +14,21 @@ from src.text import get_text
 tokenizer = AutoTokenizer.from_pretrained("gpt2")
 
 
-# def get_model_predictions(
-#     *,
-#     input_text: str,
-#     model: PreTrainedModel,
-#     tokenizer: PreTrainedTokenizer,
-# ) -> 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, :]
-#     top_3 = torch.topk(last_token, MAX_ATTEMPTS).indices.tolist()
-#     return top_3
-
-
 def make_predictions(tokenizer: PreTrainedTokenizer) -> tuple:
     """
-    Run this on startup
+    Run this on startup.
+    Returns tuple of target_prediction_pairs and target_prediction_tokens:
+
     """
     text = get_text()
     model = AutoModelForCausalLM.from_pretrained("gpt2")
     model.eval()
 
     inputs = tokenizer(text, return_tensors="pt")
-    logits = model(**inputs).logits
+
+    with torch.no_grad():
+        logits = model(**inputs).logits
+
     top_n = torch.topk(logits, MAX_ATTEMPTS)
     token_id_preds = top_n.indices.squeeze().tolist()
     tokens = list(map(tokenizer.convert_ids_to_tokens, token_id_preds))

src/reducer.py

@@ -8,6 +8,7 @@ from src.handlers import handle_out_of_attempts
 from src.handlers import handle_player_win
 from src.handlers import handle_tie
 from src.params import ReducerParams
+from src.shared import tokenizer
 from src.utils import guess_is_correct
 from src.utils import lm_is_correct
 
@@ -28,12 +29,12 @@ def _handle_guess(params: ReducerParams) -> ReducerParams:
     if params.button_label == "Next word":
         return handle_next_word(params)
 
-    if not params.prompt_text:
+    if not params.guess_field:
         return handle_no_input(params)
 
-    params.player_guesses += "\n" + params.prompt_text
-    player_correct = guess_is_correct(params.prompt_text)
-    lm_correct = lm_is_correct()
+    params.current_guesses += "\n" + params.guess_field
+    player_correct = guess_is_correct(params, tokenizer)
+    lm_correct = lm_is_correct(params)
 
     if player_correct and lm_correct:
         return handle_tie(params)

src/utils.py

@@ -2,13 +2,19 @@ import logging
 
 from transformers import PreTrainedTokenizer
 
+from src import shared
 from src.constants import MAX_ATTEMPTS
+from src.constants import STARTING_INDEX
 from src.params import ReducerParams
 from src.shared import token_id_predictions
 
 logger = logging.getLogger(__name__)
 
 
+def get_current_prompt_text(word_number):
+    return shared.tokenizer.decode(shared.all_tokens[: STARTING_INDEX + word_number])
+
+
 def get_start_and_whitespace_tokens(
     word: str,
     tokenizer: PreTrainedTokenizer,
@@ -34,13 +40,17 @@ def lm_is_correct(params: ReducerParams) -> bool:
     return current_guess == current_target
 
 
-def guess_is_correct(params: ReducerParams, current_target: int, tokenizer: PreTrainedTokenizer) -> bool:
+def guess_is_correct(params: ReducerParams, tokenizer: PreTrainedTokenizer) -> bool:
     """
     We check if the predicted token or a corresponding one with a leading whitespace
     matches that of the next token
     """
+    # FIXME: handle indexerro
+    print(STARTING_INDEX + params.word_number)
+    current_target = shared.all_tokens[STARTING_INDEX + params.word_number]
+
     logger.debug("Next token: '{}'".format(tokenizer.convert_ids_to_tokens([current_target])))
 
-    predicted_token_start, predicted_token_whitespace = get_start_and_whitespace_tokens(params.guess_field)
+    predicted_token_start, predicted_token_whitespace = get_start_and_whitespace_tokens(params.guess_field, tokenizer)
     logger.debug(tokenizer.convert_ids_to_tokens([predicted_token_start, predicted_token_whitespace]))
     return current_target in (predicted_token_start, predicted_token_whitespace)