Spaces:
Build error
Build error
import transformers | |
import numpy as np | |
import torch | |
import streamlit as st | |
from transformers import GPT2Tokenizer, GPT2LMHeadModel | |
from transformers import pipeline | |
def load_model(): | |
model_ckpt = "bankholdup/rugpt3_song_writer" | |
tokenizer = GPT2Tokenizer.from_pretrained(model_ckpt) | |
model = GPT2LMHeadModel.from_pretrained(model_ckpt) | |
return tokenizer, model | |
def set_seed(args): | |
rd = np.random.randint(100000) | |
print('seed =', rd) | |
np.random.seed(rd) | |
torch.manual_seed(rd) | |
if args.n_gpu > 0: | |
torch.cuda.manual_seed_all(rd) | |
title = st.title("Loading model") | |
tokenizer, model = load_model() | |
title.title("ruGPT3 Song Writer") | |
context = st.text_input("Введите начало песни", "Как дела? Как дела? Это новый кадиллак") | |
if st.button("Поехали", help="Может занять какое-то время"): | |
prefix_text = f"{context}" | |
encoded_prompt = tokenizer.encode(prefix_text, add_special_tokens=False, return_tensors="pt") | |
output_sequences = model.generate( | |
input_ids=encoded_prompt, | |
max_length=200 + len(encoded_prompt[0]), | |
temperature=0.95, | |
top_k=50, | |
top_p=0.95, | |
repetition_penalty=1.0, | |
do_sample=True, | |
num_return_sequences=1, | |
) | |
# Remove the batch dimension when returning multiple sequences | |
if len(output_sequences.shape) > 2: | |
output_sequences.squeeze_() | |
for generated_sequence_idx, generated_sequence in enumerate(output_sequences): | |
print("ruGPT:".format(generated_sequence_idx + 1)) | |
generated_sequence = generated_sequence.tolist() | |
# Decode text | |
text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True) | |
# Remove all text after the stop token | |
text = text[: text.find("</s>") if "</s>" else None] | |
# Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing | |
total_sequence = ( | |
prompt_text + text[len(tokenizer.decode(encoded_prompt[0], clean_up_tokenization_spaces=True)) :] | |
) | |
generated_sequences.append(total_sequence) | |
# os.system('clear') | |
st.write(total_sequence) | |