import gradio as gr
from transformers import pipeline
from huggingface_hub import login

login(token="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")

def generate_text(
    model_name,
    text,
    min_length,
    max_length
):
    models_map = {
        "Мои любимые юмореск": "gpt2-vk-aneki",
        "Калик)": "gpt2-vk-kalik",
        "бугро тред": "gpt2-vk-bugro"
    }

    model = "MesonWarrior/" + models_map[model_name]

    pipe = pipeline(
        'text-generation',
        model=model,
        tokenizer=model,
        min_length=min_length,
        max_length=max_length,
        use_auth_token=True
    )

    return pipe(text)[0]['generated_text']

def interface():
    with gr.Row():
        with gr.Column():
            with gr.Row():
                model = gr.Dropdown(
                    ["Мои любимые юморески", "Калик)", "бугро тред"],
                    label="Модель (Текст какого паблика генерировать)",
                    value="бугро тред",
                )
            text = gr.Textbox(lines=7, label="Входной текст", placeholder="Введите текст который продолжит нейросеть...")
        output = gr.Textbox(lines=12, label="Выходной текст", placeholder="Здесь будет текст сгенерированный нейросетью...")
    with gr.Row():
        with gr.Column():
            min_length = gr.Slider(
                minimum=0, maximum=128, value=32, step=1,
                label="Минимальная длина",
                info="Минимальное количество символов в выходном тексте."
            )
            max_length = gr.Slider(
                minimum=0, maximum=512, value=64, step=1,
                label="Максимальная длина",
                info="Максимальное количество символов в выходном тексте."
            )
        with gr.Column():
            with gr.Row():
                generate_btn = gr.Button(
                    "Сгенерировать", variant="primary", label="Generate",
                )

        generate_btn.click(
            fn=generate_text,
            inputs=[
                model,
                text,
                min_length,
                max_length,
            ],
            outputs=output,
        )

with gr.Blocks(
    title="GPT2 VK") as demo:
        gr.Markdown("""
        # GPT2 VK
        Файнтюны модели [ai-forever/rugpt3medium_based_on_gpt2](https://huggingface.co/ai-forever/rugpt3medium_based_on_gpt2) по вашим любимым пабликам ВКонтакте.
        #### Паблики представленные в моделях:
        - Мои любимые юморески 🎩
        - Калик) 🍏🍎💨
        - бугро тред 💥
        """)
        interface()

demo.queue().launch()