ilva / pages /gpt.py
Ilvir's picture
Update pages/gpt.py
f53026e
raw
history blame
2 kB
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import streamlit as st
import torch
import textwrap
import plotly.express as px
st.header(':green[Генерация текста 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('models/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, 10, 3)
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=8,
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()
st.write('**_Результат_** 👇')
for i, out_ in enumerate(out):
with st.expander(f'Текст {i+1}:'):
st.write(textwrap.fill(tokenizer.decode(out_), 100))
st.image("pict/wow.png")