Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,52 +1,67 @@
 import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForQuestionAnswering
-import torch
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 
-# Настройки
-MODEL_NAME = "AlexKay/xlm-roberta-large-qa-multilingual-finedtuned-ru"
-CONTEXT = """
-Абитуриенты магистратуры подают документы через личный кабинет. Требуется диплом бакалавра/специалиста. 
-Средний балл диплома рассчитывается как сумма всех оценок, деленная на количество дисциплин (без учета ВКР). 
-Доступные программы: 
-- Информатика и вычислительная техника (код 09.04.01).  
-- Управление персоналом (код 38.04.03). 
-- Физика (код 04.04.01). 
-Параллельное зачисление на две программы возможно только на платной основе. 
-Прием документов дистанционный через Госуслуги или чат-бот вуза. 
-"""
+# Заголовок приложения
+st.title("🤖 Помощник по вопросам магистратуры")
+st.write("Задайте вопрос о поступлении, обучении или программах магистратуры")
 
-# Загрузка модели
+# Инициализация модели (кешируется)
 @st.cache_resource
-def load_qa_model():
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForQuestionAnswering.from_pretrained(MODEL_NAME)
-    return pipeline("question-answering", model=model, tokenizer=tokenizer)
+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
 
-qa_pipeline = load_qa_model()
+generator = load_model()
 
-# Интерфейс
-st.title("🤖 Консультант для абитуриентов магистратуры")
-st.info("Отвечаю только на вопросы о магистратуре! Примеры: программы, документы, средний балл.")
+# Форма для ввода вопроса
+with st.form("question_form"):
+    user_input = st.text_area("Ваш вопрос:", "Что такое магистратура?")
+    submit_button = st.form_submit_button("Получить ответ")
 
-question = st.text_input("Задайте вопрос:")
-if question:
-    # Фильтр тематики
-    forbidden_keywords = ["бакалавр", "егэ", "олимпиад", "школ", "аспирантур", "специалитет"]
-    if any(word in question.lower() for word in forbidden_keywords):
-        st.error("Извините, я консультирую только по магистратуре. Задайте вопрос о программах, документах или экзаменах.")
+# Обработка вопроса
+if submit_button:
+    if not user_input.strip():
+        st.warning("Пожалуйста, введите вопрос")
     else:
-        # Поиск ответа в контексте
-        result = qa_pipeline(question=question, context=CONTEXT)
-        st.subheader("Ответ:")
-        st.write(result["answer"])
-        st.caption(f"Точность: {result['score']:.2f}")
-        
-        # Показать программы, если спрашивают о них
-        if "программ" in question.lower():
-            st.divider()
-            st.write("**Все программы магистратуры:**")
-            st.markdown("""
-            - 🖥️ **Информатика и вычислительная техника** (09.04.01)  
-            - 👥 **Управление персоналом** (38.04.03)  
-            - 🔬 **Физика** (04.04.01)
-            """)
+        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)}")
+
+# Информация о модели
+st.divider()
+st.markdown("""
+**О приложении:**
+- Использует русскоязычную модель `rugpt3small_based_on_gpt2`
+- Отвечает на вопросы о магистратуре
+- Работает на CPU (для ускорения используйте GPU)
+""")