import os
import gradio as gr

from transformers import T5ForConditionalGeneration, AutoTokenizer
# from transformers import pipeline

auth_token = os.environ.get("CLARIN_KNEXT")

model_name = "clarin-knext/plt5-large-poquad-ext-qa-autotoken"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=auth_token)
model = T5ForConditionalGeneration.from_pretrained(model_name, use_auth_token=auth_token)


default_generate_kwargs = {
            "max_length": 192,
            "num_beams": 2,
            "length_penalty": 0,
            "early_stopping": True,
        }

# keywords_pipe = pipeline(model=model, tokenizer=tokenizer, **default_generate_kwargs)

examples = [
    ["extract question: Kiedy Arek ma spotkanie dzisiaj? context: Arek ma spotkanie przez cały dzeiń. Bardzo to lubi."]]



def generate(context):
    context = context + " </s>"
    inputs = tokenizer(
        context, 
        max_length=512,
        add_special_tokens=True,
        truncation=True,
        padding=False,
        return_tensors="pt"
        )
    
    outs = model.generate(
        input_ids=inputs['input_ids'], 
        attention_mask=inputs['attention_mask'],
        **default_generate_kwargs
        )
    prediction = tokenizer.decode(outs[0], skip_special_tokens=True)
    return prediction

demo = gr.Interface(
    fn=generate,
    inputs=gr.inputs.Textbox(lines=5, label="Input Text"),
    outputs=gr.outputs.Textbox(label="Keywords"),
    examples=examples,
)

demo.launch()