File size: 4,258 Bytes
b31a2f6
 
 
1c98656
b31a2f6
3774ebc
b31a2f6
a575d46
 
b31a2f6
 
 
a575d46
 
27bdf54
 
 
b31a2f6
1c98656
90ca31e
40ec473
4f971e0
 
b31a2f6
 
 
4f971e0
 
 
 
b31a2f6
 
4f971e0
b31a2f6
4f971e0
 
b31a2f6
 
3408ba9
 
b31a2f6
1c98656
53f4560
7a6a8cc
4f971e0
 
b31a2f6
 
 
952c01c
532af44
b31a2f6
35a6465
a0b0d25
b31a2f6
 
 
 
 
 
 
b17f4f4
b31a2f6
df265bf
1c98656
df265bf
1c98656
df265bf
1c98656
 
 
90ca31e
27bdf54
b31a2f6
 
ff0ccdb
1c98656
 
 
 
35a6465
53ebfaf
431a7d1
 
 
1c98656
 
b31a2f6
1c98656
b31a2f6
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
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()