Spaces:
Sleeping
Sleeping
| 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() | |