|
|
from transformers import GPT2LMHeadModel, GPT2Tokenizer |
|
|
import streamlit as st |
|
|
import torch |
|
|
import textwrap |
|
|
|
|
|
|
|
|
st.markdown('## Генерация текста GPT-моделью') |
|
|
|
|
|
tokenizer = GPT2Tokenizer.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2') |
|
|
model = GPT2LMHeadModel.from_pretrained( |
|
|
'sberbank-ai/rugpt3small_based_on_gpt2', |
|
|
output_attentions = False, |
|
|
output_hidden_states = False, |
|
|
) |
|
|
|
|
|
model.load_state_dict(torch.load('modelgpt.pt', map_location=torch.device('cpu'))) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
length = st.sidebar.slider('Длина генерируемой последовательности:', 8, 256, 15) |
|
|
num_samples = st.sidebar.slider('Число генераций:', 1, 10, 1) |
|
|
temperature = st.sidebar.slider('Температура:', 1.0, 10.0, 2.0) |
|
|
top_k = st.sidebar.slider('Количество наиболее вероятных слов генерации:', 10, 200, 50) |
|
|
top_p = st.sidebar.slider('Минимальная суммарная вероятность топовых слов:', 0.4, 1.0, 0.9) |
|
|
|
|
|
|
|
|
prompt = st.text_input('Введите текст:') |
|
|
if st.button('Сгенерировать текст'): |
|
|
|
|
|
with torch.inference_mode(): |
|
|
prompt = tokenizer.encode(prompt, return_tensors='pt') |
|
|
out = model.generate( |
|
|
input_ids=prompt, |
|
|
max_length=length, |
|
|
num_beams=5, |
|
|
do_sample=True, |
|
|
temperature=temperature, |
|
|
top_k=top_k, |
|
|
top_p=top_p, |
|
|
no_repeat_ngram_size=3, |
|
|
num_return_sequences=num_samples, |
|
|
).cpu().numpy() |
|
|
for i, out_ in enumerate(out): |
|
|
st.write(f'Текст {i+1}:') |
|
|
st.write(textwrap.fill(tokenizer.decode(out_), 100)) |
|
|
|
