Get next token on guess

app.py

@@ -1,47 +1,52 @@
-# from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM
 from transformers import AutoTokenizer
 
-target_text = """\
-Obama was born in Honolulu, Hawaii.
-After graduating from Columbia University in 1983, he worked as a community organizer in Chicago.
-In 1988, he enrolled in Harvard Law School, where he was the first black president of the Harvard Law Review.
-After graduating, he became a civil rights attorney and an academic, teaching constitutional law at the University of Chicago Law School from 1992 to 2004.
-Turning to elective politics, he represented the 13th district in the Illinois Senate from 1997 until 2004, when he ran for the U.S.
-Senate.
-Obama received national attention in 2004 with his March Senate primary win, his well-received July Democratic National Convention keynote address, and his landslide November election to the Senate.
-In 2008, after a close primary campaign against Hillary Clinton, he was nominated by the Democratic Party for president and chose Joe Biden as his running mate.
-Obama was elected over Republican nominee John McCain in the presidential election and was inaugurated on January 20, 2009.
-Nine months later, he was named the 2009 Nobel Peace Prize laureate, a decision that drew a mixture of praise and criticism.
-"""
+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")
 
-def get_next_word(text: str) -> str:
-    inputs = tokenizer(text, return_tensors="pt")
     with torch.no_grad():
         logits = model(**inputs).logits
 
     last_token = logits[0, -1]
     top_3 = torch.topk(last_token, 3)
-    input_ids = list(inputs.input_ids.squeeze())
-    argmax = top_3.indices[0]
-    input_ids.append(argmax)
-    return "!!!", tokenizer.decode(input_ids)
 
+    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(target_text)
+            prompt_text = gr.Markdown()
         with gr.Row():
             with gr.Column():
                 guess = gr.Textbox(label="Guess!")
@@ -49,11 +54,7 @@ def build_demo():
             with gr.Column():
                 lm_guess = gr.Textbox(label="LM guess")
 
-        guess_btn.click(get_next_word, inputs=guess, outputs=[prompt_text, lm_guess], api_name="get_next_word")
-        # examples = gr.Examples(
-        #     examples=["I went to the supermarket yesterday.", "Helen is a good swimmer."], inputs=[guess]
-        # )
-
+        guess_btn.click(handle_guess, inputs=guess, outputs=[prompt_text, lm_guess], api_name="get_next_word")
     return demo
 
 
@@ -67,8 +68,8 @@ def wip_sign():
 
 
 def main():
-    # demo = build_demo()
     demo = wip_sign()
+    # demo = build_demo()
     demo.launch(debug=True)
 
 

state.py

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

text.py

@@ -0,0 +1,16 @@
+target_text = """\
+Obama was born in Honolulu, Hawaii.
+After graduating from Columbia University in 1983, he worked as a community organizer in Chicago.
+In 1988, he enrolled in Harvard Law School, where he was the first black president of the Harvard Law Review.
+After graduating, he became a civil rights attorney and an academic, teaching constitutional law at the University of Chicago Law School from 1992 to 2004.
+Turning to elective politics, he represented the 13th district in the Illinois Senate from 1997 until 2004, when he ran for the U.S.
+Senate.
+Obama received national attention in 2004 with his March Senate primary win, his well-received July Democratic National Convention keynote address, and his landslide November election to the Senate.
+In 2008, after a close primary campaign against Hillary Clinton, he was nominated by the Democratic Party for president and chose Joe Biden as his running mate.
+Obama was elected over Republican nominee John McCain in the presidential election and was inaugurated on January 20, 2009.
+Nine months later, he was named the 2009 Nobel Peace Prize laureate, a decision that drew a mixture of praise and criticism.
+"""
+
+
+def get_text():
+    return target_text