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