Create chatbot interface

.gitignore

@@ -0,0 +1 @@
+onnx

app.py

@@ -1,28 +1,39 @@
-from transformers import LongT5ForConditionalGeneration, AutoTokenizer
+from transformers import AutoTokenizer
 import time
+import gradio as gr
+from optimum.onnxruntime import ORTModelForSeq2SeqLM
+from optimum.utils import NormalizedConfigManager
+
+@classmethod
+def _new_get_normalized_config_class(cls, model_type):
+    return cls._conf["t5"]
+
+NormalizedConfigManager.get_normalized_config_class = _new_get_normalized_config_class
+
 
 N = 2 # Number of previous QA pairs to use for context
 MAX_NEW_TOKENS = 128 # Maximum number of tokens for each answer
 
-tokenizer = AutoTokenizer.from_pretrained("tryolabs/long-t5-tglobal-base-blogpost-cqa")
-model = LongT5ForConditionalGeneration.from_pretrained("tryolabs/long-t5-tglobal-base-blogpost-cqa")
+tokenizer = AutoTokenizer.from_pretrained("tryolabs/long-t5-tglobal-base-blogpost-cqa-onnx")
+model = ORTModelForSeq2SeqLM.from_pretrained("tryolabs/long-t5-tglobal-base-blogpost-cqa-onnx")
+
 
 with open("context_short.txt", "r") as f:
     context = f.read()
 
-def build_input(question, user_history=[], bot_history=[]):
+def build_input(question, state=[[],[]]):
     model_input = f"{context} || "
-    previous = min(len(bot_history[1:]), N)
+    previous = min(len(state[1][1:]), N)
     for i in range(previous, 0, -1):
-        prev_question = user_history[-i-1]
-        prev_answer = bot_history[-i]
+        prev_question = state[0][-i-1]
+        prev_answer = state[1][-i]
         model_input += f"<Q{i}> {prev_question} <A{i}> {prev_answer} "
     model_input += f"<Q> {question} <A> "
     return model_input
 
-def get_model_answer(question, user_history=[], bot_history=[]):
+def get_model_answer(question, state=[[],[]]):
     start = time.perf_counter()
-    model_input = build_input(question, user_history, bot_history)
+    model_input = build_input(question, state)
     end = time.perf_counter()
     print(f"Build input: {end-start}")
     start = time.perf_counter()
@@ -34,11 +45,21 @@ def get_model_answer(question, user_history=[], bot_history=[]):
     end = time.perf_counter()
     print(f"Tokenize: {end-start}")
     start = time.perf_counter()
-    encoded_output = model.generate(input_ids=input_ids, attention_mask=attention_mask, do_sample=True, max_new_tokens=MAX_NEW_TOKENS)
+    encoded_output = model.generate(input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=MAX_NEW_TOKENS)
     answer = tokenizer.decode(encoded_output[0], skip_special_tokens=True)
     end = time.perf_counter()
     print(f"Generate: {end-start}")
-    user_history.append(question)
-    bot_history.append(answer)
-    return answer, user_history, bot_history
+    state[0].append(question)
+    state[1].append(answer)
+    responses = [(state[0][i], state[1][i]) for i in range(len(state[0]))]
+    return responses, state
+
+with gr.Blocks() as demo:
+    state = gr.State([[],[]])
+    chatbot = gr.Chatbot()
+    text = gr.Textbox(label="Ask a question (press enter to submit)", default_value="How are you?")
+
+    text.submit(get_model_answer, [text, state], [chatbot, state])
+    text.submit(lambda x: "", text, text)
 
+demo.launch()

requirements.txt

@@ -1,2 +1,8 @@
 transformers
-torch
+torch
+onnx==1.12.0
+onnxconverter-common==1.13.0
+onnxruntime==1.13.1
+onnxruntime-tools==1.7.0
+openvino==2022.2.0
+optimum @ git+https://github.com/huggingface/optimum.git@4c3b1c14f07c8e3780d9c9765b3992a90fab3349