import transformers import torch import tokenizers import gradio as gr import re from PIL import Image def get_model_gpt(model_name,tokenizer_name): tokenizer = transformers.GPT2Tokenizer.from_pretrained(tokenizer_name) model = transformers.GPT2LMHeadModel.from_pretrained(model_name) model.eval() 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) model.eval() return model, tokenizer def predict_gpt(text, model, tokenizer, temperature=1.0): input_ids = tokenizer.encode(text, return_tensors="pt") 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].split('Описание :')[1] return 'Описание : '+ generated_text def predict_t5(text, model, tokenizer, temperature=1.2): input_ids = tokenizer.encode(text, return_tensors="pt") with torch.no_grad(): out = model.generate(input_ids, do_sample=True, num_beams=4, temperature=temperature, top_p=0.35, max_length=1024, length_penalty = 5.5, 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,text): if model == 'GPT': result = predict_gpt(text, gpt_model, gpt_tokenizer) elif model == 'T5': result = predict_t5(text, t5_model, t5_tokenizer) return result gpt_model, gpt_tokenizer = get_model_gpt('mipatov/rugpt3_nb_descr', 'mipatov/rugpt3_nb_descr') t5_model, t5_tokenizer = get_model_t5('mipatov/rut5_nb_descr', 'mipatov/rut5_nb_descr') example = ' Название : Super_NB 2001 Gaming;\n Диагональ экрана : 21 " ;\n Разрешение экрана : 1337x228 ;\n Поверхность экрана : матовая ;\n Тип матрицы : nfc ;\n Процессор : CMD processor 7 ядер 16.16 ГГц (46.0 ГГц, в режиме Turbo) ;\n Оперативная память : 28 Гб DDR5 ;\n Тип видеокарты : интегрированный ;\n Видеокарта : qwerty-grafics ;\n SSD : 720 Гб ;\n Wi-Fi : ДА, 802.11 a/b/g/n/ac ;\n Bluetooth : ДА, v5.0 ;\n Кабельная сеть : 10/100/1000 (Gigabit Ethernet) Мбит/с ;\n USB 2.0 : 13 ;\n USB 3.0 : 22 ;\n HDMI : 11 ;\n Операционная система : CMD-shell ;\n Веб-камера : встроенная ;\n Микрофон : есть ;\n Разъем наушники/микрофон : комбинированный разъем ;\n Акустическая система : стереодинамики ;\n Цвет клавиатуры : черный ;\n Цифровой блок клавиатуры : есть ;\n Подсветка клавиш клавиатуры : есть ;\n Тип батареи : Al-Ion ;\n Количество ячеек батареи : 36 cell ;\n Энергоемкость батареи : 176 Wh ;\n Цвет : черный ;\n Размеры : 1.23 х 456 х 78.9 мм ;\n Вес : 19 кг ;\n Гарантия : 322 мес. ;\n Материал корпуса : пластик ;\n Время работы от батареи : 82ч ;\n Кард-ридер : есть WA SD ;' demo = gr.Interface( fn=generate, inputs=[ gr.components.Dropdown(label="Модель", choices=('GPT', 'T5')), gr.components.Textbox(label="Характеристики",value = example), ], outputs=[ gr.components.Textbox(label="Описание",interactive = False), ], ) demo.launch()