Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,67 +1,77 @@
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import pipeline
+import time
 
-# Заголовок приложения
-st.title("🤖 Помощник по вопросам магистратуры")
-st.write("Задайте вопрос о поступлении, обучении или программах магистратуры")
+# Настройки для экономии памяти
+MAX_LENGTH = 200
+MODEL_NAME = "sberbank-ai/rugpt3small_based_on_gpt2"
 
-# Инициализация модели (кешируется)
-@st.cache_resource
+st.set_page_config(page_title="Магистратура Помощник")
+st.title("🎓 Помощник по вопросам магистратуры")
+st.write("Задайте вопрос о поступлении или обучении в магистратуре")
+
+# Кешируем загрузку модели
+@st.cache_resource(show_spinner=False)
 def load_model():
-    model_name = "sberbank-ai/rugpt3small_based_on_gpt2"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-    generator = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        device=-1  # CPU (для GPU измените на 0)
-    )
-    return generator
+    with st.spinner("Загружаем модель... Это займет около 30 секунд"):
+        return pipeline(
+            "text-generation",
+            model=MODEL_NAME,
+            device_map="auto"
+        )
 
 generator = load_model()
 
-# Форма для ввода вопроса
+# История диалога
+if "history" not in st.session_state:
+    st.session_state.history = []
+
+# Функция генерации ответа
+def generate_answer(question):
+    prompt = f"Вопрос о магистратуре: {question}\nОтвет:"
+    try:
+        result = generator(
+            prompt,
+            max_length=MAX_LENGTH,
+            num_return_sequences=1,
+            temperature=0.7,
+            repetition_penalty=1.2,
+            pad_token_id=50256
+        )
+        return result[0]['generated_text'].split("Ответ:")[-1].strip()
+    except Exception as e:
+        return f"Ошибка: {str(e)}"
+
+# Форма ввода
 with st.form("question_form"):
-    user_input = st.text_area("Ваш вопрос:", "Что такое магистратура?")
-    submit_button = st.form_submit_button("Получить ответ")
+    user_input = st.text_input("Ваш вопрос:", placeholder="Что нужно для поступления?")
+    submit_button = st.form_submit_button("Спросить")
 
 # Обработка вопроса
-if submit_button:
-    if not user_input.strip():
-        st.warning("Пожалуйста, введите вопрос")
-    else:
-        with st.spinner("Генерируем ответ..."):
-            # Форматируем промпт для лучшего ответа
-            prompt = f"Вопрос о магистратуре: {user_input}\nОтвет:"
-            
-            # Генерация ответа
-            try:
-                results = generator(
-                    prompt,
-                    max_length=300,
-                    num_return_sequences=1,
-                    temperature=0.7,
-                    repetition_penalty=1.5,
-                    pad_token_id=50256
-                )
-                
-                # Извлекаем и очищаем ответ
-                answer = results[0]['generated_text'].split("Ответ:")[-1].strip()
-                answer = answer.split('\n')[0]  # Берем первый абзац
-                
-                # Выводим результат
-                st.subheader("Ответ:")
-                st.write(answer)
-                
-            except Exception as e:
-                st.error(f"Ошибка генерации: {str(e)}")
+if submit_button and user_input:
+    with st.spinner("Генерируем ответ..."):
+        start_time = time.time()
+        response = generate_answer(user_input)
+        st.session_state.history.append((user_input, response))
+        
+        # Очищаем поле ввода после отправки
+        st.rerun()
+
+# Показываем историю диалога
+for i, (question, answer) in enumerate(st.session_state.history[::-1]):
+    st.divider()
+    st.subheader(f"❓ {question}")
+    st.write(f"💡 {answer}")
 
 # Информация о модели
-st.divider()
-st.markdown("""
-**О приложении:**
-- Использует русскоязычную модель `rugpt3small_based_on_gpt2`
-- Отвечает на вопросы о магистратуре
-- Работает на CPU (для ускорения используйте GPU)
+st.sidebar.markdown("""
+### О приложении
+Использует ру��скоязычную модель [rugpt3small](https://huggingface.co/sberbank-ai/rugpt3small_based_on_gpt2) для ответов на вопросы о магистратуре.
+
+**Примеры вопросов:**
+- Какие документы нужны для поступления?
+- Сколько длится обучение?
+- Чем магистратура отличается от бакалавриата?
+- Какие есть направления?
+- Есть ли бюджетные места?
 """)