import transformers
import torch
import tokenizers
import gradio as gr
import re


def get_model_gpt(model_name,tokenizer_name):
    tokenizer = transformers.GPT2Tokenizer.from_pretrained(tokenizer_name)
    model = transformers.GPT2LMHeadModel.from_pretrained(model_name)
    return model, tokenizer

def get_model_t5(model_name,tokenizer_name):
    tokenizer = transformers.T5Tokenizer.from_pretrained(tokenizer_name)
    model = transformers.T5ForConditionalGeneration.from_pretrained(model_name)
    return model, tokenizer


def predict_gpt(text, model, tokenizer, temperature=1.0):
    text = text.replace('\n','')+' Описание:'
    input_ids = tokenizer.encode(text, return_tensors="pt")

    model.eval()
    with torch.no_grad():
        out = model.generate(input_ids, 
                            do_sample=True,
                            num_beams=3,
                            temperature=temperature,
                            top_p=0.75,
                            max_length=1024,
                            eos_token_id = tokenizer.eos_token_id,
                            pad_token_id = tokenizer.pad_token_id,
                            repetition_penalty = 2.5,
                            num_return_sequences = 1,
                             output_attentions = True,
                             return_dict_in_generate=True,
                            )
    decode = lambda x : tokenizer.decode(x, skip_special_tokens=True)
    generated_text = list(map(decode, out['sequences']))[0].replace(text,'')
    return generated_text

def predict_t5(text, model, tokenizer, temperature=1.2):
    text = text.replace('\n','')
    input_ids = tokenizer.encode(text, return_tensors="pt")
    
    model.eval()
    with torch.no_grad():
        out = model.generate(input_ids, 
                        do_sample=True,
                        num_beams=4,
                        temperature=temperature,
                        top_p=0.35,
                        max_length=512,
                        length_penalty = -1.0,
                        output_attentions = True,
                        return_dict_in_generate=True,
                        repetition_penalty = 2.5,
                        num_return_sequences = 1
                        )
    decode = lambda x : tokenizer.decode(x, skip_special_tokens=True)
    generated_text = list(map(decode, out['sequences']))[0]
    return 'Описание :'+generated_text

def generate(model,temp,text):
    if model == 'GPT':
                result = predict_gpt(text, gpt_model, gpt_tokenizer,temp)
    elif model == 'T5':
                result = predict_t5(text, t5_model, t5_tokenizer,temp)
    return result
         

gpt_model, gpt_tokenizer = get_model_gpt('gpt/', 'gpt/')
t5_model, t5_tokenizer = get_model_t5('mipatov/rut5_nb_descr', 'mipatov/rut5_nb_descr')


example = '  Название : Super_NB 2001 Gaming;\n Диагональ экрана : 25" ;\n Разрешение экрана : 1337x228 ;\n Поверхность экрана : матовая ;\n Тип матрицы : nfc ;\n Процессор : CPU-007 7 ядер 16.16 ГГц (46.0 ГГц, в режиме Turbo) ;\n Оперативная память : 13 Гб DDR5 ;\n Тип видеокарты : интегрированный ;\n Видеокарта : qwerty-grafics ;\n SSD : 720 Гб ;\n Операционная система : MY_OS ;\n Тип батареи : Li-Ion ;\n Количество ячеек батареи : 36 cell ;\n Энергоемкость батареи : 176 Wh ;\n Цвет : черный ;\n Размеры : 1.23 х  456 х  78.9 мм ;\n Вес : 19 кг ;\n Гарантия : 322 мес. ;\n Время работы от батареи : 82ч ;\n Кард-ридер : есть WA SD ;'
demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.components.Dropdown(label="Модель", choices=('GPT', 'T5')),
        gr.components.Slider(label="Вариативность",value = 1.2,minimum = 1.0,maximum = 3.0,step = 0.1),
        gr.components.Textbox(label="Характеристики",value = example),
    ],
    outputs=[
        gr.components.Textbox(label="Описание",interactive = False),
    ],
)

demo.launch()