Update app.py

app.py

@@ -4,7 +4,6 @@ import numpy as np
 import nltk
 from nltk.tokenize import sent_tokenize, word_tokenize
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from nltk.data import find
 import functools
 
 
@@ -17,6 +16,9 @@ st.set_page_config(
 
 @functools.lru_cache(maxsize=None)
 def get_model():
+    """
+    Кэшируемая функция для загрузки модели и токенизатора.
+    """
     # Загрузка модели
     model = AutoModelForCausalLM.from_pretrained('model')
     # Загрузка токенизатора
@@ -24,39 +26,23 @@ 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')
-    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)
-
-
 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)
@@ -66,18 +52,27 @@ def process_reviews(reviews):
 
 
 def load_nltk_data():
+    """
+    Функция для загрузки данных NLTK.
+    Проверяет, загружены ли данные, и при необходимости загружает их.
+    """
     try:
-        find('tokenizers/punkt')
-        find('tokenizers/punkt_tab')
+        nltk.data.find('tokenizers/punkt')
         print("Данные уже загружены.")
     except LookupError:
         print("Загрузка данных NLTK...")
-        nltk.download(['punkt', 'punkt_tab'])
-
+        nltk.download('punkt')
 
+        
 def preprocess_input(input_text):
+    """
+    Функция для предварительной обработки входного текста.
+    Убирает лишние символы, исправляет предложения.
+    """
     input_text = input_text.split(":")[-1].strip()
     sentences = sent_tokenize(input_text)
+    if len(sentences) == 0:
+        return ""
     last_sentence = sentences[-1]
     if not last_sentence.endswith('.'):
         sentences.pop()
@@ -93,11 +88,14 @@ def preprocess_input(input_text):
 
 
 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,
+        max_length=200,
         num_return_sequences=1,
         no_repeat_ngram_size=2,
         do_sample=True,
@@ -110,22 +108,28 @@ def generate_review(input_text):
 
 
 def main():
+    """
+    Основная функция приложения Streamlit.
+    """
     if 'btn_predict' not in st.session_state:
         st.session_state['btn_predict'] = False
 
+    # Ввод данных пользователем
     category = st.text_input("Категория:", value="Кондитерская")
     rating = st.slider("Рейтинг", 1, 5, 1)
     key_words = st.text_input("Ключевые слова", value="десерт, торт, цена")
 
-    # Ввод новых параметров
+    # Формируем входной текст
     input_text = f"Категория: {category}; Рейтинг: {rating}; Ключевые слова: {key_words} -> Отзыв:"
 
-    if st.button('Generate'):
+    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)
 
+
 if __name__ == "__main__":
+    load_nltk_data()
     main()