Update app.py

app.py

@@ -5,6 +5,7 @@ import nltk
 from nltk.tokenize import sent_tokenize, word_tokenize
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from nltk.data import find
+import functools
 
 
 # Настройка конфигурации страницы Streamlit
@@ -14,21 +15,7 @@ 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()
+@functools.lru_cache(maxsize=None)
 def get_model():
     # Загрузка модели
     model = AutoModelForCausalLM.from_pretrained('model')
@@ -37,7 +24,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')
@@ -55,36 +42,73 @@ def gen_review(input_text):
     return tokenizer.decode(output[0], skip_special_tokens=True)
 
 
-def capitalize_and_punctuate(text):
-    download_nltk_data()
-    # Разделяем текст на предложения
-    sentences = sent_tokenize(text)
+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():
+    try:
+        find('tokenizers/punkt')
+        find('tokenizers/punkt_tab')
+        print("Данные уже загружены.")
+    except LookupError:
+        print("Загрузка данных NLTK...")
+        nltk.download(['punkt', 'punkt_tab'])
+
+
+def preprocess_input(input_text):
+    input_text = input_text.split(":")[-1].strip()
+    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
 
 
-# Главная функция
+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)
+
+
 def main():
     if 'btn_predict' not in st.session_state:
         st.session_state['btn_predict'] = False
@@ -98,11 +122,10 @@ def main():
 
     if st.button('Generate'):
         with st.spinner('Генерация отзыва...'):
-            generated_text = gen_review(input_text)
-            generated_text = capitalize_and_punctuate(generated_text)        
-        st.success("Готово!")
-        st.text(generated_text)
-
+            processed_input = preprocess_input(input_text)
+            generated_text = generate_review(processed_input)
+            st.success("Готово!")
+            st.text(generated_text)
 
 if __name__ == "__main__":
     main()