Update pages/GPT.py

pages/GPT.py

@@ -5,6 +5,7 @@ import transformers
 import random
 import textwrap
 
+@st.cache
 def load_model():
     model_finetuned = transformers.AutoModelWithLMHead.from_pretrained(
         'tinkoff-ai/ruDialoGPT-small',
@@ -17,6 +18,7 @@ def load_model():
 
 def preprocess_text(text_input, tokenizer):
     prompt = tokenizer.encode(text_input, return_tensors='pt')
+    return prompt
 
 def predict_sentiment(model, prompt, temp, num_generate):
     with torch.inference_mode():
@@ -31,6 +33,7 @@ def predict_sentiment(model, prompt, temp, num_generate):
             no_repeat_ngram_size=3,
             num_return_sequences=num_generate,
             ).cpu().numpy()
+    print(result)
     return result
 
 st.title('Text generation with dreambook')
@@ -38,8 +41,11 @@ st.title('Text generation with dreambook')
 model, tokenizer = load_model()
 
 text_input = st.text_input("Enter some text about movie")
+print(text_input)
 max_len = st.slider('Length of sequence', 0, 500, 250)
+print(max_len)
 temp = st.slider('Temperature', 1, 30, 1)
+print(temp)
 if st.button('Generate a random number of sequences'):
     num_generate = random.randint(1,5)
     st.write(f'Number of sequences: {num_generate}')