haha tokens go brrr

app.py

@@ -1,24 +1,46 @@
 import gradio as gr
-import random
-import time
+from threading import Thread
+from transformers import AutoModelForCausalLM, AutoTokenizer, IteratorStreamer
 
+
+# Global variable loading
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+print("Loading the model...")
+model = AutoModelForCausalLM.from_pretrained("gpt2")
+print("Done!")
+
+
+# Gradio app
 with gr.Blocks() as demo:
+    def user(user_message, history):
+        return "", history + [[user_message, None]]
+
     chatbot = gr.Chatbot()
     msg = gr.Textbox()
     clear = gr.Button("Clear")
 
-    def user(user_message, history):
-        return "", history + [[user_message, None]]
+    def update_chatbot(history):
+        user_query = history[-1][0]
+        history[-1][1] = ""
+        model_inputs = tokenizer([user_query], return_tensors="pt")
+
+        # Start generation on a separate thread, so that we don't block the UI. The text is pulled from the streamer
+        # in the main thread.
+        streamer = IteratorStreamer(tokenizer)
+        generate_kwargs = dict(model_inputs, streamer=streamer, max_new_tokens=200, do_sample=True)
+        t = Thread(target=model.generate, kwargs=generate_kwargs)
+        t.start()
 
-    def bot(history):
-        bot_message = random.choice(["Yes", "No"])
-        history[-1][1] = bot_message
-        time.sleep(1)
+        # Pull the generated text from the streamer, and update the chatbot.
+        for new_text in streamer:
+            history[-1][1] += new_text
+            yield history
         return history
 
-    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot, chatbot, chatbot
+    msg.submit(user, [msg, chatbot], [msg, chatbot]).then(
+        update_chatbot, chatbot, chatbot
     )
-    clear.click(lambda: None, None, chatbot, queue=False)
+    clear.click(lambda: None, None, chatbot)
 
+demo.queue()
 demo.launch()

requirements.txt

@@ -1 +1 @@
-git+https://github.com/huggingface/transformers.git  # transformers from `main`
+git+https://github.com/gante/transformers.git@streamer_iterator  # transformers from dev branch