Update app.py

app.py

@@ -7,7 +7,6 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 st.set_page_config(page_title="Генератор отзывов", initial_sidebar_state="expanded")
 st.title("Генератор отзывов на основе ИИ")
 st.write("Создавайте текстовые отзывы с помощью нейросети на основе категорий, рейтинга и ключевых слов.")
-st.sidebar.title("Параметры генерации")
 
 # ================ Загрузка модели ================
 @st.cache_resource
@@ -25,6 +24,17 @@ def load_model():
         st.error(f"Ошибка при загрузке модели: {e}")
         st.stop()
 
+# ================ Предустановленные стили ================
+def get_preset_params(style):
+    """Возвращает параметры генерации для выбранного стиля."""
+    if style == "Строгий":
+        return {"temperature": 0.3, "top_p": 0.8, "top_k": 30, "no_repeat_ngram_size": 3}
+    elif style == "Умеренный":
+        return {"temperature": 0.7, "top_p": 0.9, "top_k": 50, "no_repeat_ngram_size": 2}
+    elif style == "Безумный":
+        return {"temperature": 1.5, "top_p": 1.0, "top_k": 100, "no_repeat_ngram_size": 1}
+    return {}
+
 # ================ Форматирование текста ================
 def format_text(text):
     """Форматирует текст: заглавные буквы и корректные точки."""
@@ -40,48 +50,29 @@ def format_text(text):
             formatted_sentences.append(sentence)
     return " ".join(formatted_sentences)
 
-# ================ Оптимизированная генерация ================
-def generate_text(model, tokenizer, prompt, params):
-    """Генерирует текст блоками для ускорения."""
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
-    output_ids = input_ids
-
-    for _ in range(params["max_length"] // 10):  # Генерация по 10 токенов
-        with torch.no_grad():
-            outputs = model.generate(
-                input_ids=output_ids,
-                max_new_tokens=10,
-                temperature=params["temperature"],
-                top_p=params["top_p"],
-                top_k=params["top_k"],
-                do_sample=params["do_sample"],
-                no_repeat_ngram_size=params["no_repeat_ngram_size"],
-                pad_token_id=tokenizer.eos_token_id,
-            )
-            new_tokens = outputs[0][output_ids.shape[1]:]
-            output_ids = outputs
-            new_text = tokenizer.decode(new_tokens, skip_special_tokens=True)
-            yield new_text.strip()
-
 # ================ Основное приложение ================
 def main():
     model, tokenizer = load_model()
 
-    # Параметры генерации
-    params = {
-        "max_length": st.sidebar.slider("Максимальная длина", 50, 300, 150, help="Максимальная длина текста."),
-        "temperature": st.sidebar.slider("Температура", 0.01, 2.0, 0.7, 0.05, help="Степень случайности."),
-        "top_p": st.sidebar.slider("Top-p", 0.01, 1.0, 0.9, 0.05, help="Фильтрация по вероятности."),
-        "top_k": st.sidebar.number_input("Top-k", 1, 100, 50, help="Ограничение по количеству кандидатов."),
-        "do_sample": st.sidebar.checkbox("Случайная выборка", value=True, help="Включает случайность генерации."),
-        "no_repeat_ngram_size": st.sidebar.number_input("Размер n-грамм", 1, 10, 2, help="Предотвращает повторение фраз."),
-    }
-
-    # Поля ввода
+    # Основные параметры
     category = st.selectbox("Категория:", ["Кафе", "Ресторан", "Парк", "Музей"], help="Выберите категорию.")
     rating = st.slider("Рейтинг", 1, 5, 5)
     key_words = st.text_input("Ключевые слова", "вкусно, уютно, быстро", help="Введите ключевые слова для отзыва.")
 
+    # Выбор стиля генерации
+    style = st.selectbox("Стиль генерации:", ["Строгий", "Умеренный", "Безумный"], help="Выберите стиль текста.")
+    params = get_preset_params(style)
+
+    # Расширенные настройки в сворачивающемся блоке
+    with st.expander("Дополнительные настройки"):
+        params["temperature"] = st.slider("Температура", 0.01, 2.0, params["temperature"], 0.05, help="Степень случайности.")
+        params["top_p"] = st.slider("Top-p", 0.01, 1.0, params["top_p"], 0.05, help="Фильтрация по вероятности.")
+        params["top_k"] = st.number_input("Top-k", 1, 100, params["top_k"], help="Ограничение по количеству кандидатов.")
+        params["no_repeat_ngram_size"] = st.number_input("Размер n-грамм", 1, 10, params["no_repeat_ngram_size"], help="Предотвращает повторение фраз.")
+        max_length = st.slider("Максимальная длина", 50, 300, 150, help="Максимальная длина текста.")
+        num_return_sequences = st.number_input("Количество вариантов", 1, 5, 1, help="Сколько вариантов текста нужно сгенерировать.")
+
+    # Генерация текста
     if st.button("Сгенерировать"):
         if not key_words:
             st.warning("Введите ключевые слова.")
@@ -90,16 +81,31 @@ def main():
         input_prompt = f"Напиши отзыв для заведения. Категория: {category}. Рейтинг: {rating}. Ключевые слова: {key_words}."
         st.info("Генерация текста...")
 
-        placeholder = st.empty()
-        full_text = ""
+        placeholders = [st.empty() for _ in range(num_return_sequences)]
+        texts = []
 
-        for chunk in generate_text(model, tokenizer, input_prompt, params):
-            full_text += " " + chunk
-            formatted_text = format_text(full_text)
-            placeholder.write(formatted_text)
+        for i in range(num_return_sequences):
+            with torch.no_grad():
+                output = model.generate(
+                    tokenizer.encode(input_prompt, return_tensors="pt"),
+                    max_new_tokens=max_length,
+                    temperature=params["temperature"],
+                    top_p=params["top_p"],
+                    top_k=params["top_k"],
+                    do_sample=True,
+                    no_repeat_ngram_size=params["no_repeat_ngram_size"],
+                    pad_token_id=tokenizer.eos_token_id,
+                )
+            generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+            formatted_text = format_text(generated_text)
+            texts.append(formatted_text)
+            placeholders[i].write(f"**Вариант {i + 1}:**\n{formatted_text}")
 
         st.success("Генерация завершена!")
-        st.download_button("Скачать отзыв", formatted_text, "generated_review.txt")
+
+        # Скачивание всех вариантов
+        all_texts = "\n\n".join([f"Вариант {i + 1}:\n{text}" for i, text in enumerate(texts)])
+        st.download_button("Скачать все отзывы", all_texts, "generated_reviews.txt")
 
 if __name__ == "__main__":
     main()