import gradio as gr

import generation.generators as generators
from datamanager import models
from train import train

demo = gr.Blocks(
    title="Text Generator v2"
)
info_text = """
## English
**Language**: English
**Quality**: 9/10
**Sources**: ChatGPT, https://pastebin.com/WYvij310
## English-XL (Temporarily unavailable)
**Language**: English
**Quality**: Broken
**Sources**: https://pastebin.com/WYvij310, https://whiletrue.neocities.org/lte, https://rainbowfluffysheep.wordpress.com/the-longest-text-ever/
## Russian
**Language**: Russian
**Quality**: 7/10
**Sources**: https://goroda.murman.ru/index.php?topic=6508.20;wap2
## Russian V2
**Language**: Russian
**Quality**: 7/10
**Sources**: http://staging.budsvetom.com/literature_items/ochen-dlinnyy-tekst
## Russian War and Piece
**Language**: Russian
**Quality**: 6/10
**Sources**: https://gist.github.com/Semionn/bdcb66640cc070450817686f6c818897
"""


def train_model(model_name, progress=gr.Progress(track_tqdm=True)):
    train(generators.find_model(model_name))
    return 'Trained ' + model_name

with demo:
    gr.Markdown('# Text Generator V2')

    model = gr.Dropdown([model_info['name']for model_info in models.values()], label="Model", value="English")

    with gr.Row():
        with gr.Column(scale=2):

            with gr.Tab('Generate'):
                inp = gr.Textbox(label="Prompt")
                word_count = gr.Number(40, precision=1, label="Max word count")
                stop_chance = gr.Slider(0, 1, 0.5, step=0.05, label='Stop probability')
                btn = gr.Button("Submit", variant="primary")

            with gr.Tab('Continue'):
                inp2 = gr.Textbox(label="Text")
                word_count2 = gr.Number(40, precision=1, label="Max word count")
                btn2 = gr.Button("Submit", variant="primary")

            with gr.Tab('Train'):
                train_out = gr.Textbox(label="Training result")
                train_btn = gr.Button("Submit", variant="primary")

        out = gr.Textbox(label="Output")

    btn.click(generators.generate, [inp, word_count, model, stop_chance], out, api_name="generate")
    btn2.click(generators.cont, [inp2, word_count2, model], out, api_name="continue")
    train_btn.click(train_model, model, train_out, api_name="train")

    with gr.Accordion(label="Information", open=False):
        gr.Markdown(info_text)
    
    gr.HTML('<footer style="text-align:center;opacity:0.5"> v2.2 · Quality update · Data version: 3</footer>')

demo.queue()
demo.launch()