app.py

from transformers import LongT5ForConditionalGeneration, AutoTokenizer
import time

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")

with open("context_short.txt", "r") as f:
 context = f.read()

def build_input(question, user_history=[], bot_history=[]):
 model_input = f"{context} || "
 previous = min(len(bot_history[1:]), N)
 for i in range(previous, 0, -1):
 prev_question = user_history[-i-1]
 prev_answer = bot_history[-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=[]):
 start = time.perf_counter()
 model_input = build_input(question, user_history, bot_history)
 end = time.perf_counter()
 print(f"Build input: {end-start}")
 start = time.perf_counter()
 encoded_inputs = tokenizer(model_input, max_length=3000, truncation=True, return_tensors="pt")
 input_ids, attention_mask = (
 encoded_inputs.input_ids,
 encoded_inputs.attention_mask
 )
 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)
 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