Update app.py

app.py

@@ -11,6 +11,7 @@ st.write("Создавайте текстовые отзывы с помощью
 # ================ Загрузка модели ================
 @st.cache_resource
 def load_model():
+    """Загрузка модели и токенизатора (кэшируется для предотвращения перезагрузки)."""
     model_path = "./model"
     if not os.path.exists(model_path):
         st.error(f"Модель не найдена по пути {model_path}. Проверьте путь.")
@@ -26,86 +27,116 @@ def load_model():
 
 # ================ Рейтинг звёздочками ================
 def star_rating():
+    """Выбор рейтинга с отображением звёздочек."""
     stars = ["⭐", "⭐⭐", "⭐⭐⭐", "⭐⭐⭐⭐", "⭐⭐⭐⭐⭐"]
-    return st.radio("Рейтинг:", range(1, 6), format_func=lambda x: stars[x - 1])
-
-# ================ Стиль генерации ================
-def get_style_params(style):
-    styles = {
-        "Строгий": {"temperature": 0.3, "top_p": 0.8, "top_k": 30},
-        "Умеренный": {"temperature": 0.7, "top_p": 0.9, "top_k": 50},
-        "Безумный": {"temperature": 1.2, "top_p": 1.0, "top_k": 100},
-    }
-    return styles.get(style, {"temperature": 0.7, "top_p": 0.9, "top_k": 50})
+    selected_star = st.radio("Рейтинг:", range(1, 6), format_func=lambda x: stars[x - 1], help="Выберите оценку от 1 до 5 звёзд.")
+    return selected_star
 
 # ================ Форматирование текста ================
 def format_text(text):
+    """Форматирует текст: заглавные буквы и точки."""
     sentences = text.split(". ")
     formatted = []
     for sentence in sentences:
         sentence = sentence.strip()
         if sentence:
-            sentence = sentence.capitalize() if not sentence[0].isupper() else sentence
-            sentence += "." if not sentence.endswith((".", "!", "?")) else ""
+            if not sentence[0].isupper():
+                sentence = sentence.capitalize()
+            if not sentence.endswith((".", "!", "?")):
+                sentence += "."
             formatted.append(sentence)
     return " ".join(formatted)
 
-# ================ Генерация текста ================
-def generate_text(model, tokenizer, prompt, params, max_length):
+# ================ Потоковая генерация текста ================
+def stream_generate(model, tokenizer, prompt, params, max_length, batch_size):
+    """Генерирует текст блоками с обновлением контекста."""
     input_ids = tokenizer.encode(prompt, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model.generate(
-            input_ids=input_ids,
-            max_length=max_length,
-            temperature=params["temperature"],
-            top_p=params["top_p"],
-            top_k=params["top_k"],
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    output_ids = input_ids
+
+    for _ in range(max_length // batch_size):
+        if st.session_state.get("stop_generation", False):
+            st.warning("Генерация остановлена!")
+            break
+
+        with torch.no_grad():
+            outputs = model.generate(
+                input_ids=output_ids,
+                max_new_tokens=batch_size,
+                temperature=params["temperature"],
+                top_p=params["top_p"],
+                top_k=params["top_k"],
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id,
+                no_repeat_ngram_size=params["no_repeat_ngram_size"],
+            )
+            new_tokens = outputs[0, output_ids.shape[1]:]
+            output_ids = torch.cat([output_ids, new_tokens.unsqueeze(0)], dim=-1)  # Обновляем контекст
+            new_text = tokenizer.decode(new_tokens, skip_special_tokens=True)
+            yield new_text.strip()
+
+            if tokenizer.eos_token_id in new_tokens:
+                break
 
 # ================ Основное приложение ================
 def main():
-    model, tokenizer = load_model()
+    # Инициализация состояния для кнопки остановки
+    if "stop_generation" not in st.session_state:
+        st.session_state.stop_generation = False
 
-    # Инициализация состояния
-    if "is_generating" not in st.session_state:
-        st.session_state.is_generating = False  # По умолчанию генерация выключена
+    model, tokenizer = load_model()
 
     # Основные параметры
     category = st.selectbox("Категория:", ["Кафе", "Ресторан", "Парк", "Музей"])
     rating = star_rating()
-    key_words = st.text_input("Ключевые слова", "вкусно, уютно, быстро")
+    key_words = st.text_input("Ключевые слова", "вкусно, уютно, быстро", help="Введите ключевые слова для использования в отзыве.")
 
-    # Стиль генерации
-    style = st.selectbox("Стиль генерации:", ["Строгий", "Умеренный", "Безумный"])
-    style_params = get_style_params(style)
+    # Стили генерации
+    style = st.selectbox("Стиль генерации:", ["Строгий", "Умеренный", "Безумный"],
+                         help="Выберите стиль текста: Строгий – формальный и точный; Безумный – творческий и креативный.")
+    style_params = {"Строгий": (0.3, 30), "Умеренный": (0.7, 50), "Безумный": (1.5, 100)}
+    temperature, top_k = style_params[style]
 
     # Дополнительные настройки
     with st.expander("Дополнительные настройки"):
         max_length = st.slider("Максимальная длина", 50, 300, 150)
+        num_variants = st.number_input("Количество вариантов", 1, 5, 1)
+        batch_size = st.slider("Размер батча токенов", 1, 20, 5, help="Количество токенов за шаг.")
+        temperature = st.slider("Температура", 0.01, 2.0, temperature, 0.05)
+        top_p = st.slider("Top-p", 0.01, 1.0, 0.9, 0.05)
+        top_k = st.number_input("Top-k", 1, 100, top_k)
+        no_repeat_ngram_size = st.slider("Размер n-грамм для предотвращения повторов", 1, 10, 3)
 
-    # Управление кнопкой
-    button_label = "🔴 Остановить" if st.session_state.is_generating else "🟢 Сгенерировать"
-
-    if st.button(button_label):
-        st.session_state.is_generating = not st.session_state.is_generating  # Переключение состояния
+    # Кнопка для остановки генерации
+    if st.button("Остановить генерацию"):
+        st.session_state.stop_generation = True
 
     # Генерация текста
-    if st.session_state.is_generating:
+    if st.button("Сгенерировать"):
+        st.session_state.stop_generation = False  # Сбрасываем состояние остановки
+        if not key_words:
+            st.warning("Введите ключевые слова.")
+            return
+
+        input_prompt = f"Напиши отзыв для заведения. Категория: {category}. Рейтинг: {rating} из 5. Ключевые слова: {key_words}."
         st.info("Генерация текстов...")
-        try:
-            prompt = f"Категория: {category}. Рейтинг: {rating}. Ключевые слова: {key_words}. Стиль: {style}."
-            generated_text = generate_text(
-                model, tokenizer, prompt, style_params, max_length
-            )
-            st.success("Генерация завершена!")
-            st.write(format_text(generated_text))
-        except Exception as e:
-            st.error(f"Ошибка во время генерации: {e}")
-        finally:
-            st.session_state.is_generating = False  # Возврат в исходное состояние
+
+        placeholders = [st.empty() for _ in range(num_variants)]
+        texts = [""] * num_variants
+
+        for i in range(num_variants):
+            for chunk in stream_generate(
+                model, tokenizer, input_prompt, 
+                {"temperature": temperature, "top_p": top_p, "top_k": top_k, "no_repeat_ngram_size": no_repeat_ngram_size},
+                max_length, batch_size
+            ):
+                texts[i] += chunk + " "
+                placeholders[i].write(f"**Вариант {i + 1}:**\n{format_text(texts[i])}")
+
+        st.success("Генерация завершена!")
+
+        # Скачивание всех вариантов
+        all_texts = "\n\n".join([f"Вариант {i + 1}:\n{format_text(text)}" for i, text in enumerate(texts)])
+        st.download_button("Скачать все отзывы", all_texts, "generated_reviews.txt")
 
 if __name__ == "__main__":
     main()