Update app.py

app.py

@@ -1,99 +1,103 @@
-import streamlit as st
-import pandas as pd
-import numpy as np
-import nltk
-from nltk.tokenize import sent_tokenize, word_tokenize
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-# Setting up the page configuration for Streamlit App
-st.set_page_config(
-    page_title="Generate reviews",
-    # layout="wide",
-    initial_sidebar_state="expanded"
-)
-
-
-# Загрузка модели и токенизатора
-#@st.cache_data()
-def get_model():
-    # Загрузка модели
-    model = AutoModelForCausalLM.from_pretrained('model')
-    # Загрузка токенизатора
-    tokenizer = AutoTokenizer.from_pretrained('model')
-    return (model, tokenizer)
-
-
-# Генерация отзыва
-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 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
-
-
-# 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)
-        generated_text = capitalize_and_punctuate(generated_text)
-        st.text(generated_text)
-
-
-
-if __name__ == "__main__":
-    main()
-
+import streamlit as st
+import pandas as pd
+import numpy as np
+import nltk
+from nltk.tokenize import sent_tokenize, word_tokenize
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Setting up the page configuration for Streamlit App
+st.set_page_config(
+    page_title="Generate reviews",
+    # layout="wide",
+    initial_sidebar_state="expanded"
+)
+
+
+# Загрузка модели и токенизатора
+#@st.cache_data()
+def get_model():
+    # Загрузка модели
+    model = AutoModelForCausalLM.from_pretrained('model')
+    # Загрузка токенизатора
+    tokenizer = AutoTokenizer.from_pretrained('model')
+    return (model, tokenizer)
+
+
+# Генерация отзыва
+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 capitalize_and_punctuate(text):
+    nltk.download('punkt')
+    # Разделяем текст на предложения
+    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
+
+
+# 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...'):
+            generated_text = capitalize_and_punctuate(generated_text)
+            st.text(generated_text)
+        st.success("Done!")
+        
+
+
+
+if __name__ == "__main__":
+    main()
+