Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,62 +1,52 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForQuestionAnswering
+import torch
 
-# Настройка заголовка
-st.title("🎓 Ассистент по магистратуре")
-st.write("Задайте вопросы о поступлении, программах или требованиях")
-
-# Инициализация модели (кэшируется для ускорения)
-@st.cache_resource
-def load_model():
-    return pipeline(
-        "text-generation", 
-        model="IlyaGusev/saiga_llama3_8b",
-        device_map="auto"
-    )
+# Настройки
+MODEL_NAME = "AlexKay/xlm-roberta-large-qa-multilingual-finedtuned-ru"
+CONTEXT = """
+Абитуриенты магистратуры подают документы через личный кабинет. Требуется диплом бакалавра/специалиста. 
+Средний балл диплома рассчитывается как сумма всех оценок, деленная на количество дисциплин (без учета ВКР). 
+Доступные программы: 
+- Информатика и вычислительная техника (код 09.04.01).  
+- Управление персоналом (код 38.04.03). 
+- Физика (код 04.04.01). 
+Параллельное зачисление на две программы возможно только на платной основе. 
+Прием документов дистанционный через Госуслуги или чат-бот вуза. 
+"""
 
 # Загрузка модели
-generator = load_model()
+@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)
 
-# История диалога
-if "messages" not in st.session_state:
-    st.session_state.messages = [
-        {"role": "assistant", "content": "Привет! Я помогу с вопросами о магистратуре. Спрашивайте!"}
-    ]
+qa_pipeline = load_qa_model()
 
-# Отображение истории
-for msg in st.session_state.messages:
-    st.chat_message(msg["role"]).write(msg["content"])
+# Интерфейс
+st.title("🤖 Консультант для абитуриентов магистратуры")
+st.info("Отвечаю только на вопросы о магистратуре! Примеры: программы, документы, средний балл.")
 
-# Обработка ввода
-prompt = st.chat_input("Ваш вопрос...")
-if prompt:
-    # Добавляем вопрос в историю
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    st.chat_message("user").write(prompt)
-    
-    # Формируем контекст диалога
-    context = "\n".join(
-        f"{'Ты:' if m['role']=='user' else 'Ассистент:'} {m['content']}" 
-        for m in st.session_state.messages
-    )
-    
-    # Генерация ответа
-    with st.spinner("Думаю..."):
-        try:
-            response = generator(
-                context,
-                max_new_tokens=512,
-                temperature=0.7,
-                top_p=0.9,
-                do_sample=True
-            )[0]['generated_text']
-            
-            # Извлекаем только последний ответ
-            assistant_reply = response.split("Ассистент:")[-1].strip()
-            
-        except Exception as e:
-            assistant_reply = f"Ошибка: {str(e)}"
-    
-    # Добавляем ответ в историю
-    st.session_state.messages.append({"role": "assistant", "content": assistant_reply})
-    st.chat_message("assistant").write(assistant_reply)
+question = st.text_input("Задайте вопрос:")
+if question:
+    # Фильтр тематики
+    forbidden_keywords = ["бакалавр", "егэ", "олимпиад", "школ", "аспирантур", "специалитет"]
+    if any(word in question.lower() for word in forbidden_keywords):
+        st.error("Извините, я консультирую только по магистратуре. Задайте вопрос о программах, документах или экзаменах.")
+    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)
+            """)