Update app.py

app.py

@@ -2,30 +2,31 @@ import streamlit as st
 import pandas as pd
 import numpy as np
 import nltk
+nltk.download('punkt')
 from nltk.tokenize import sent_tokenize, word_tokenize
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Setting up the page configuration for Streamlit App
+
+# Настройка конфигурации страницы Streamlit
 st.set_page_config(
     page_title="Generate reviews",
-    # layout="wide",
     initial_sidebar_state="expanded"
 )
 
 
 # Загрузка модели и токенизатора
-#@st.cache_data()
+@st.cache_resource()
 def get_model():
     # Загрузка модели
     model = AutoModelForCausalLM.from_pretrained('model')
     # Загрузка токенизатора
     tokenizer = AutoTokenizer.from_pretrained('model')
-    return (model, tokenizer)
+    return model, tokenizer
 
 
 # Генерация отзыва
-def gen_review(input_text):  
-    (model, tokenizer) = get_model()
+def gen_review(input_text):
+    model, tokenizer = get_model()
     input_ids = tokenizer.encode(input_text, return_tensors='pt')
     output = model.generate(
         input_ids,
@@ -42,62 +43,52 @@ def gen_review(input_text):
 
 
 def capitalize_and_punctuate(text):
-    nltk.download('punkt')
     # Разделяем текст на предложения
     sentences = sent_tokenize(text)
-    
+
     # Проверка последнего предложения
     last_sentence = sentences[-1]
-    if not last_sentence.endswith('.'):  
+    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
 
 
-# Main function
+# Главная функция
 def main():
     if 'btn_predict' not in st.session_state:
         st.session_state['btn_predict'] = False
 
-    # Sidebar
-    # st.sidebar.markdown(''' # New York City Taxi Trip Duration''')
-    # st.sidebar.image("img/taxi_img.png")
-    
     category = st.text_input("Категория:", value="Кондитерская")
     rating = st.slider("Рейтинг", 1, 5, 1)
     key_words = st.text_input("Ключевые слова", value="десерт, торт, цена")
-    
+
     # Ввод новых параметров
     input_text = f"Категория: {category}; Рейтинг: {rating}; Ключевые слова: {key_words} -> Отзыв:"
 
-
-    st.session_state['btn_predict'] = st.button('Generate')
-    if st.session_state['btn_predict']:
-        generated_text = gen_review(input_text)
-        with st.spinner('Wait for it...'):
+    if st.button('Generate'):
+        with st.spinner('Генерация отзыва...'):
+            generated_text = gen_review(input_text)
             generated_text = capitalize_and_punctuate(generated_text)
             st.text(generated_text)
-        st.success("Done!")
-        
-
+        st.success("Готово!")
 
 
 if __name__ == "__main__":
     main()
-