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
-import functools
 
 
 # Настройка конфигурации страницы Streamlit
@@ -14,11 +13,14 @@ st.set_page_config(
 )
 
 
-@functools.lru_cache(maxsize=None)
+def download_nltk_data():
+    nltk.download('punkt')
+    nltk.download('punkt_tab')
+
+
+# Загрузка модели и токенизатора
+@st.cache_data()
 def get_model():
-    """
-    Кэшируемая функция для загрузки модели и токенизатора.
-    """
     # Загрузка модели
     model = AutoModelForCausalLM.from_pretrained('model')
     # Загрузка токенизатора
@@ -26,91 +28,13 @@ def get_model():
     return model, tokenizer
 
 
-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 load_nltk_data():
-    """
-    Функция для загрузки данных NLTK.
-    """
-    nltk.download('punkt')    
-    nltk.download('punkt_tab')
-
-        
-def preprocess_input(input_text):
-    """
-    Функция для предварительной обработки входного текста.
-    Убирает лишние символы, исправляет предложения.
-    """
-    # Удаляем лишние символы и извлекаем текст после двоеточия
-    print(f" ДО {input_text}")
-    input_text = input_text.split(":")[-1].strip()
-    print(f" ПОСЛЕ {input_text}")
-    
-    # Токенизация предложений
-    sentences = sent_tokenize(input_text)
-    
-    # Проверка на наличие предложений
-    if len(sentences) == 0:
-        return ""
-    
-    # Удаляем последнее предложение, если оно не заканчивается на точку
-    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)
-    
-    # Добавление точки в конце, если финальный текст не пустой
-    if final_text and not final_text.endswith('.'):
-        final_text += '.'
-    
-    return final_text
-
-
-def generate_review(input_text):
-    """
-    Функция для генерации отзыва на основе входного текста.
-    """
+# Генерация отзыва
+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=200,
+        max_length=300,
         num_return_sequences=1,
         no_repeat_ngram_size=2,
         do_sample=True,
@@ -122,29 +46,54 @@ def generate_review(input_text):
     return tokenizer.decode(output[0], skip_special_tokens=True)
 
 
+def capitalize_and_punctuate(text):    
+    # Разделяем текст на предложения
+    sentences = sent_tokenize(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
+
+
+# Главная функция
 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)
+            generated_text = gen_review(input_text)
+            generated_text = capitalize_and_punctuate(generated_text)        
+        st.success("Готово!")
+        st.text(generated_text)
 
 
 if __name__ == "__main__":
-    load_nltk_data()
+    download_nltk_data()
     main()