import gradio as gr
from transformers import GPT2Tokenizer
from thirdai import bolt, licensing

licensing.activate("7511CC-0E24D7-69439D-5D6CBA-33AAFD-V3")

tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = bolt.GenerativeModel.load("./generative.model")


def generate(prompt):
    prompt = tokenizer.encode(prompt)

    stream = model.streaming_generation(
        input_tokens=prompt,
        prediction_chunk_size=2,
        max_predictions=80,
        beam_width=10,
    )

    for res in stream:
        yield tokenizer.decode(res)


with gr.Blocks() as demo:
    output = gr.TextArea(label="Output")
    prompt = gr.Textbox(
        label="Prompt",
    )
    prompt.submit(generate, inputs=[prompt], outputs=[output])

    btn = gr.Button(value="Generate")
    btn.click(generate, inputs=[prompt], outputs=[output])

    gr.ClearButton(components=[prompt, output])

if __name__ == "__main__":
    demo.queue()
    demo.launch()