Update app.py

app.py

@@ -4,8 +4,7 @@ import numpy as np
 import nltk
 from nltk.tokenize import sent_tokenize, word_tokenize
 from transformers import AutoModelForCausalLM, AutoTokenizer
-nltk.download('punkt')
-nltk.download('punkt_tab')
+from nltk.data import find
 
 
 # Настройка конфигурации страницы Streamlit
@@ -15,8 +14,21 @@ st.set_page_config(
 )
 
 
+def download_nltk_data():
+    try:
+        # Проверяем, установлены ли данные
+        find('tokenizers/punkt')
+        find('tokenizers/punkt_tab')
+        print("Данные уже загружены.")
+    except LookupError:
+        # Если данные не найдены, загружаем их
+        print("Загрузка данных NLTK...")
+        nltk.download('punkt')
+        nltk.download('punkt_tab')
+
+
 # Загрузка модели и токенизатора
-# @st.cache_data()
+@st.cache_data()
 def get_model():
     # Загрузка модели
     model = AutoModelForCausalLM.from_pretrained('model')
@@ -44,6 +56,7 @@ def gen_review(input_text):
 
 
 def capitalize_and_punctuate(text):
+    download_nltk_data()
     # Разделяем текст на предложения
     sentences = sent_tokenize(text)
 
@@ -84,11 +97,12 @@ def main():
     input_text = f"Категория: {category}; Рейтинг: {rating}; Ключевые слова: {key_words} -> Отзыв:"
 
     if st.button('Generate'):
-        #with st.spinner('Генерация отзыва...'):
-        generated_text = gen_review(input_text)
-        generated_text = capitalize_and_punctuate(generated_text)        
-        #st.success("Готово!")
+        with st.spinner('Генерация отзыва...'):
+            generated_text = gen_review(input_text)
+            generated_text = capitalize_and_punctuate(generated_text)        
+        st.success("Готово!")
         st.text(generated_text)
 
+
 if __name__ == "__main__":
     main()