import ctranslate2
from transformers import AutoTokenizer

import gradio as gr



def generate_prompt(history):
    prompt = ""
    for chain in history[:-1]:
        prompt += f"<human>: {chain[0]}\n<bot>: {chain[1]}\n"
    prompt += f"<human>: {history[-1][0]}\n<bot>:"
    tokens = tokenizer.convert_ids_to_tokens(tokenizer.encode(prompt))
    return tokens

def generate(question):
    history = [[question, ""]]
    tokens = generate_prompt(history)

    results = translator.translate_batch(
        [tokens],
        beam_size=1,
        max_decoding_length = 256,
        repetition_penalty = 1.8,
    )
    answer = tokenizer.convert_tokens_to_string(results[0].hypotheses[0])
    return answer



translator = ctranslate2.Translator("model")
tokenizer = AutoTokenizer.from_pretrained("DKYoon/mt5-xl-lm-adapt")
end_token = "</s>"
end_token_id = tokenizer.encode(end_token)[0]


demo = gr.Interface(fn=generate, description="Space by @theodotus. Source: https://huggingface.co/spaces/theodotus/pythia-uk", inputs="text", outputs="text")
demo.queue(1)
if __name__ == "__main__":
    demo.launch()