import os
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

# tokenizer = AutoTokenizer.from_pretrained(
#     "milyiyo/paraphraser-german-mt5-small", use_auth_token=os.environ["AUTH_TOKEN"])
# model = AutoModelForSeq2SeqLM.from_pretrained(
#     "milyiyo/paraphraser-german-mt5-small", use_auth_token=os.environ["AUTH_TOKEN"])

tokenizer = AutoTokenizer.from_pretrained("milyiyo/paraphraser-german-mt5-small")
model = AutoModelForSeq2SeqLM.from_pretrained("milyiyo/paraphraser-german-mt5-small")


def paraphrase(sentence: str, count: str):
    p_count = int(count)
    if p_count <= 0 or len(sentence.strip()) == 0:
        return {'result': []}
    sentence_input = sentence
    text = f"paraphrase: {sentence_input} </s>"
    encoding = tokenizer.encode_plus(text, padding=True, return_tensors="pt")
    input_ids, attention_masks = encoding["input_ids"], encoding["attention_mask"]
    outputs = model.generate(
        input_ids=input_ids, attention_mask=attention_masks,
        max_length=512,  # 256,
        do_sample=True,
        top_k=120,
        top_p=0.95,
        early_stopping=True,
        num_return_sequences=p_count
    )
    res = []
    for output in outputs:
        line = tokenizer.decode(
            output, skip_special_tokens=True, clean_up_tokenization_spaces=True)
        res.append(line)
    return {'result': res}
        
def paraphrase_dummy(sentence: str, count: str):
    return {'result': []}


iface = gr.Interface(fn=paraphrase,
                     inputs=[
                         gr.inputs.Textbox(lines=2, placeholder=None, label='Sentence'),
                         gr.inputs.Number(default=3, label='Paraphrases count'),
                     ],
                     outputs=[gr.outputs.JSON(label=None)])
iface.launch()