Update pages/📖History_Mystery.py

pages/📖History_Mystery.py

@@ -9,33 +9,38 @@ st.title("""
  History Mystery
  """)
 # Добавление слайдера
-temperature = st.slider("Градус дичи", 1.0, 20.0, 1.0)
-max_length = st.slider(" Длина сгенерированного отрывка", 40, 120, 2)
+temperature = st.slider("Градус дичи", 1, 20, 1)
+max_length = st.slider(" Длина сгенерированного отрывка", 60, 120, 2)
 # Загрузка модели и токенизатора
 # model = GPT2LMHeadModel.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2')
 # tokenizer = GPT2Tokenizer.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2')
 # #Задаем класс модели (уже в streamlit/tg_bot)
-model_GPT = GPT2LMHeadModel.from_pretrained(
-     'sberbank-ai/rugpt3small_based_on_gpt2',
-     output_attentions = False,
-     output_hidden_states = False,
- )
-tokenizer_GPT = GPT2Tokenizer.from_pretrained(
+
+@st.cache
+def load_gpt():
+    model_GPT = GPT2LMHeadModel.from_pretrained(
      'sberbank-ai/rugpt3small_based_on_gpt2',
      output_attentions = False,
      output_hidden_states = False,
- )
+    )
+    tokenizer_GPT = GPT2Tokenizer.from_pretrained(
+        'sberbank-ai/rugpt3small_based_on_gpt2',
+        output_attentions = False,
+        output_hidden_states = False,
+        )
+    model_GPT.load_state_dict(torch.load('model_history.pt', map_location=torch.device('cpu')))
+    return model_GPT, tokenizer_GPT
 
 # # Вешаем сохраненные веса на нашу модель
-model_GPT.load_state_dict(torch.load('model_history.pt', map_location=torch.device('cpu')))
+
 # Функция для генерации текста
-def generate_text(prompt):
+def generate_text(model_GPT, tokenizer_GPT, prompt):
     # Преобразование входной строки в токены
     input_ids = tokenizer_GPT.encode(prompt, return_tensors='pt')
 
     # Генерация текста
     output = model_GPT.generate(input_ids=input_ids, max_length=70, num_beams=5, do_sample=True,
-                            temperature=5, top_k=50, top_p=0.6, no_repeat_ngram_size=3,
+                            temperature=1., top_k=50, top_p=0.6, no_repeat_ngram_size=3,
                             num_return_sequences=3)
 
     # Декодирование сгенерированного текста
@@ -45,6 +50,7 @@ def generate_text(prompt):
 
 # Streamlit приложение
 def main():
+    model_GPT, tokenizer_GPT = load_gpt()
     st.write("""
     # GPT-3 генерация текста
     """)
@@ -59,7 +65,7 @@ def main():
     # Обработка события нажатия кнопки
     if generate_button:
     # Вывод сгенерированного текста
-        generated_text = generate_text(prompt)
+        generated_text = generate_text(model_GPT, tokenizer_GPT, prompt)
         st.subheader("Продолжение:")
         st.write(generated_text)