| import streamlit as st | |
| import torch | |
| import numpy as np | |
| import transformers | |
| import random | |
| def load_model(): | |
| model_finetuned = transformers.AutoModelWithLMHead.from_pretrained( | |
| 'tinkoff-ai/ruDialoGPT-small', | |
| output_attentions = False, | |
| output_hidden_states = False | |
| ) | |
| model_finetuned.load_state_dict(torch.load('GPT_sonnik_only.pt', map_location=torch.device('cpu'))) | |
| tokenizer = transformers.AutoTokenizer.from_pretrained('tinkoff-ai/ruDialoGPT-small') | |
| return model_finetuned, tokenizer | |
| def preprocess_text(text_input, tokenizer): | |
| prompt = tokenizer.encode(text_input, return_tensors='pt') | |
| def predict_sentiment(model, prompt, temp, num_generate): | |
| result = model.generate( | |
| input_ids=prompt, | |
| max_length=150, | |
| num_beams=5, | |
| do_sample=True, | |
| temperature=float(temp), | |
| top_k=50, | |
| top_p=0.6, | |
| no_repeat_ngram_size=3, | |
| num_return_sequences=num_generate, | |
| ).cpu().numpy() | |
| return result | |
| st.title('Text generation with dreambook') | |
| model, tokenizer = load_model() | |
| text_input = st.text_input("Enter some text about movie") | |
| max_len = st.slider('Length of sequence', 0, 500, 250) | |
| temp = st.slider('Temperature', 0, 30, 0) | |
| if st.button('Generate a random number of sequences'): | |
| num_generate = random.randint(1,5) | |
| st.write(f'Number of sequences: {num_generate}') | |
| else: | |
| num_generate = st.text_input("Enter number of sequences") | |
| if text_input: | |
| prompt = preprocess_text(text_input, tokenizer) | |
| result = predict_sentiment(model, prompt, max_len, temp, num_generate) | |
| for i in result: | |
| st.write(textwrap.fill(tokenizer.decode(i), max_len)) | |
| if st.button('Next'): | |
| try: | |
| continue | |
| except: | |
| st.write('All sequences are return. Generate new') | |
| else: | |
| None | |