Spaces:
Sleeping
Sleeping
import numpy | |
import streamlit as st | |
import torch | |
map_location=torch.device('cpu') | |
st.title( 'Text generation by GPT model' ) | |
st.subheader( 'This application shows the difference in text generation between a rugpt3small model trained on general documents and the same model trained on jokes' ) | |
device = 'cpu' if torch.cuda.is_available() else 'cpu' | |
# Load the model tokenizer | |
from transformers import GPT2Tokenizer | |
tokenizer = GPT2Tokenizer.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2') | |
from transformers import GPT2LMHeadModel | |
# This model is simply loaded | |
model_init = GPT2LMHeadModel.from_pretrained( | |
'sberbank-ai/rugpt3small_based_on_gpt2', | |
output_attentions = False, | |
output_hidden_states = False, | |
) | |
model_init.to(device); | |
# This is a trained model, we load weights into it | |
model = GPT2LMHeadModel.from_pretrained( | |
'sberbank-ai/rugpt3small_based_on_gpt2', | |
output_attentions = False, | |
output_hidden_states = False, | |
) | |
# Load the model's saved weights and load them into the model | |
m = torch.load('model.pt', map_location=torch.device('cpu')) | |
model.load_state_dict(m) | |
model.to(device); | |
str = st.text_input( 'Enter 1-4 words of the beginning of the text, and wait a minute' , '' ) | |
# model without additional training | |
# prompt is a string that will be accepted as input and continued by the model | |
# tokenize string | |
prompt = tokenizer.encode(str, return_tensors='pt').to(device) | |
# out will contain the generation results as a list | |
out1 = model_init.generate( | |
# input string | |
input_ids=prompt, | |
# maximum length of generated sequence | |
max_length=150, | |
# num_beams | |
num_beams=5, | |
# apply sampling | |
do_sample=True, | |
# apply temperature | |
temperature=1., | |
# top words by probability | |
top_k=50, | |
# top words by total probability | |
top_p=0.6, | |
# how much (try) not to repeat n_gram in a row | |
no_repeat_ngram_size=3, | |
# how many generations to return | |
num_return_sequences=1, | |
).cpu().numpy() #).numpy() | |
st.write('\n------------------\n') | |
st.subheader( 'Texts on the model trained by documents of all subjects:' ) | |
# out contains results | |
# decoding and printing | |
n = 0 | |
for out_ in out1: | |
n += 1 | |
st.write(tokenizer.decode(out_).rpartition('.')[0],'.') | |
st.write('\n------------------\n') | |
# print(tokenizer.decode(out_)) | |
# retrained model | |
with torch.inference_mode(): | |
# prompt = 'Man asks the waiter' | |
# prompt = tokenizer.encode(str, return_tensors='pt') | |
out2 = model.generate( | |
input_ids=prompt, | |
max_length=150, | |
num_beams=1, | |
do_sample=True, | |
temperature=1., | |
top_k=5, | |
top_p=0.6, | |
no_repeat_ngram_size=2, | |
num_return_sequences=3, | |
).cpu().numpy() #).cpu().numpy() |