Update app.py

app.py

@@ -5,7 +5,6 @@ import nltk
 from nltk.tokenize import sent_tokenize, word_tokenize
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from nltk.data import find
-import functools
 
 
 # Настройка конфигурации страницы Streamlit
@@ -15,7 +14,21 @@ st.set_page_config(
 )
 
 
-@functools.lru_cache(maxsize=None)
+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()
 def get_model():
     # Загрузка модели
     model = AutoModelForCausalLM.from_pretrained('model')
@@ -24,7 +37,7 @@ def get_model():
     return model, tokenizer
 
 
-@functools.lru_cache(maxsize=None)
+# Генерация отзыва
 def gen_review(input_text):
     model, tokenizer = get_model()
     input_ids = tokenizer.encode(input_text, return_tensors='pt')
@@ -42,76 +55,39 @@ def gen_review(input_text):
     return tokenizer.decode(output[0], skip_special_tokens=True)
 
 
-def correct_sentence(sentence):
-    """Функция для исправления предложений."""
-    words = word_tokenize(sentence)
-    
-    # Делаем первую букву первого слова заглавной
-    if len(words) > 0:
-        words[0] = words[0].capitalize()
-    
-    # Собираем обратно предложение
-    corrected_sentence = ' '.join(words)
-    return corrected_sentence
-
-
-def process_reviews(reviews):
-    """Функция для обработки списка отзывов."""
-    corrected_reviews = []
-    for review in reviews:
-        sentences = sent_tokenize(review)
-        corrected_sentences = [correct_sentence(sentence) for sentence in sentences]
-        corrected_reviews.append(' '.join(corrected_sentences))
-    return corrected_reviews
-
+def capitalize_and_punctuate(text):
+    download_nltk_data()
+    # Разделяем текст на предложения
+    sentences = sent_tokenize(text)
 
-def load_nltk_data():
-    try:
-        find('tokenizers/punkt')
-        find('tokenizers/punkt_tab')
-        print("Данные уже загружены.")
-    except LookupError:
-        print("Загрузка данных NLTK...")
-        nltk.download(['punkt', 'punkt_tab'])
-
-
-def preprocess_input(input_text):
-    sentences = sent_tokenize(input_text)
+    # Проверка последнего предложения
     last_sentence = sentences[-1]
     if not last_sentence.endswith('.'):
         sentences.pop()
+
+    # Обрабатываем оставшиеся предложения
     corrected_sentences = []
     for sentence in sentences:
         words = word_tokenize(sentence)
+
+        # Делаем первую букву первого слова заглавной
         if len(words) > 0:
             words[0] = words[0].capitalize()
+
+        # Собираем обратно предложение
         corrected_sentence = ' '.join(words)
         corrected_sentences.append(corrected_sentence)
-    final_text = ' '.join(corrected_sentences)
-    return final_text
 
+    # Объединяем все предложения в единый текст
+    final_text = ' '.join(corrected_sentences)
 
-def generate_review(input_text):
-    model, tokenizer = get_model()
-    input_ids = tokenizer.encode(input_text, return_tensors='pt')
-    output = model.generate(
-        input_ids,
-        max_length=300,
-        num_return_sequences=1,
-        no_repeat_ngram_size=2,
-        do_sample=True,
-        top_p=0.95,
-        top_k=60,
-        temperature=0.9,
-        eos_token_id=tokenizer.eos_token_id,
-    )
-    return tokenizer.decode(output[0], skip_special_tokens=True)
+    return final_text
 
 
+# Главная функция
 def main():
     if 'btn_predict' not in st.session_state:
         st.session_state['btn_predict'] = False
-        load_nltk_data()
 
     category = st.text_input("Категория:", value="Кондитерская")
     rating = st.slider("Рейтинг", 1, 5, 1)
@@ -122,10 +98,11 @@ def main():
 
     if st.button('Generate'):
         with st.spinner('Генерация отзыва...'):
-            processed_input = preprocess_input(input_text)
-            generated_text = generate_review(processed_input)
-            st.success("Готово!")
-            st.text(generated_text)
+            generated_text = gen_review(input_text)
+            generated_text = capitalize_and_punctuate(generated_text)        
+        st.success("Готово!")
+        st.text(generated_text)
+
 
 if __name__ == "__main__":
     main()