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/keywords-plt5-small-shuffle"
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": 32,
            "num_beams": 4,
            "length_penalty": 1.5,
            "no_repeat_ngram_size": 3,
            "early_stopping": True,
        }

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

examples = [
    ["Prezydent Andrzej Duda odwołał Jakuba Kumocha ze stanowiska sekretarza stanu - szefa Biura Polityki Międzynarodowej i odznaczył go - za wybitne zasługi w służbie dyplomatycznej - Krzyżem Kawalerskim Orderu Odrodzenia Polski."],
]



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