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-onnx")
model = ORTModelForSeq2SeqLM.from_pretrained("tryolabs/long-t5-tglobal-base-blogpost-cqa-onnx")


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

def build_input(question, state=[[],[]]):
    model_input = f"{context} || "
    previous = min(len(state[1][1:]), N)
    for i in range(previous, 0, -1):
        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, state=[[],[]]):
    start = time.perf_counter()
    model_input = build_input(question, state)
    end = time.perf_counter()
    print(f"Build input: {end-start}")
    start = time.perf_counter()
    encoded_inputs = tokenizer(model_input, max_length=7000, 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, 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}")
    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?")
    gr.Examples(
        ["What's the name of the dataset that was built?", "what task does it focus on?", "what is that task about?"],
        text
    )

    text.submit(get_model_answer, [text, state], [chatbot, state])
    text.submit(lambda x: "", text, text)

demo.launch()