Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,140 +1,26 @@
 import streamlit as st
-from langchain_community.vectorstores import FAISS
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.document_loaders import PyPDFLoader, Docx2txtLoader, TextLoader
 from transformers import pipeline
-import os
 
-# Конфигурация
-DATA_DIR = "data"
-INDEX_DIR = "faiss_index"
-MODEL_NAME = "IlyaGusev/saiga_llama3_8b"
+# Настройка заголовка
+st.title("🎓 Ассистент по магистратуре")
+st.write("Задайте вопросы о поступлении, программах или требованиях")
 
-# Инициализация модели
+# Инициализация модели (кэшируется для ускорения)
 @st.cache_resource
-def load_llm():
+def load_model():
     return pipeline(
-        "text-generation",
-        model=MODEL_NAME,
-        device_map="auto",
-        model_kwargs={"torch_dtype": "auto"}
+        "text-generation", 
+        model="IlyaGusev/saiga_llama3_8b",
+        device_map="auto"
     )
 
-# Инициализация эмбеддингов
-@st.cache_resource
-def load_embeddings():
-    return HuggingFaceEmbeddings(model_name="cointegrated/LaBSE-en-ru")
-
-# Загрузка и обработка документов
-def process_documents():
-    documents = []
-    
-    for filename in os.listdir(DATA_DIR):
-        filepath = os.path.join(DATA_DIR, filename)
-        try:
-            if filename.endswith(".pdf"):
-                loader = PyPDFLoader(filepath)
-            elif filename.endswith(".docx"):
-                loader = Docx2txtLoader(filepath)
-            elif filename.endswith(".txt"):
-                loader = TextLoader(filepath)
-            else:
-                continue
-                
-            documents.extend(loader.load())
-        except Exception as e:
-            st.error(f"Ошибка загрузки {filename}: {str(e)}")
-    
-    if not documents:
-        return None
-    
-    # Разделение текста на чанки
-    text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=500,
-        chunk_overlap=100
-    )
-    chunks = text_splitter.split_documents(documents)
-    
-    # Создание векторного хранилища
-    embeddings = load_embeddings()
-    vectorstore = FAISS.from_documents(chunks, embeddings)
-    vectorstore.save_local(INDEX_DIR)
-    
-    return vectorstore
-
-# Поиск релевантных документов
-def retrieve_docs(query):
-    if os.path.exists(INDEX_DIR):
-        embeddings = load_embeddings()
-        vectorstore = FAISS.load_local(INDEX_DIR, embeddings)
-    else:
-        vectorstore = process_documents()
-        if vectorstore is None:
-            return []
-    
-    results = vectorstore.similarity_search(query, k=3)
-    return [doc.page_content for doc in results]
-
-# Генерация ответа с RAG
-def generate_with_rag(query, history):
-    # Получаем релевантные документы
-    context_docs = retrieve_docs(query)
-    
-    if not context_docs:
-        context = "Информация не найдена в документах."
-    else:
-        context = "\n\n".join([f"[Документ {i+1}]: {doc}" for i, doc in enumerate(context_docs)])
-    
-    # Формируем промпт
-    system_prompt = """
-    Ты ассистент по вопросам магистратуры. Отвечай ТОЛЬКО на основе предоставленной информации.
-    Если в контексте нет ответа - скажи "Я не нашел информации по этому вопросу в документах".
-    """
-    
-    history_str = "\n".join([
-        f"{'Студент' if msg['role']=='user' else 'Ассистент'}: {msg['content']}" 
-        for msg in history
-    ])
-    
-    full_prompt = f"""
-    <|system|>{system_prompt}</s>
-    <|context|>
-    {context}
-    </s>
-    <|history|>
-    {history_str}
-    </s>
-    <|user|>{query}</s>
-    <|assistant|>
-    """
-    
-    # Генерируем ответ
-    generator = load_llm()
-    response = generator(
-        full_prompt,
-        max_new_tokens=1024,
-        temperature=0.3,
-        do_sample=True,
-        eos_token_id=128001
-    )
-    
-    return response[0]['generated_text'].split("<|assistant|>")[-1].strip()
-
-# Интерфейс Streamlit
-st.title("🎓 Ассистент по магистратуре с RAG")
-st.write("Загрузите документы в папку 'data' и задавайте вопросы")
-
-# Загрузка документов
-if st.sidebar.button("Обновить базу знаний"):
-    with st.spinner("Обработка документов..."):
-        process_documents()
-    st.sidebar.success("База знаний обновлена!")
+# Загрузка модели
+generator = load_model()
 
 # История диалога
 if "messages" not in st.session_state:
     st.session_state.messages = [
-        {"role": "assistant", "content": "Привет! Задайте вопрос о магистратуре, и я отвечу на основе документов."}
+        {"role": "assistant", "content": "Привет! Я помогу с вопросами о магистратуре. Спрашивайте!"}
     ]
 
 # Отображение истории
@@ -142,28 +28,35 @@ for msg in st.session_state.messages:
     st.chat_message(msg["role"]).write(msg["content"])
 
 # Обработка ввода
-if prompt := st.chat_input("Ваш вопрос о магистратуре..."):
+prompt = st.chat_input("Ваш вопрос...")
+if prompt:
     # Добавляем вопрос в историю
     st.session_state.messages.append({"role": "user", "content": prompt})
     st.chat_message("user").write(prompt)
     
-    # Генерация ответа с RAG
-    with st.spinner("Ищу информацию..."):
+    # Формируем контекст диалога
+    context = "\n".join(
+        f"{'Ты:' if m['role']=='user' else 'Ассистент:'} {m['content']}" 
+        for m in st.session_state.messages
+    )
+    
+    # Генерация ответа
+    with st.spinner("Думаю..."):
         try:
-            response = generate_with_rag(
-                prompt, 
-                st.session_state.messages[-5:]  # Последние 5 сообщений как контекст
-            )
+            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:
-            response = f"Ошибка: {str(e)}"
+            assistant_reply = f"Ошибка: {str(e)}"
     
     # Добавляем ответ в историю
-    st.session_state.messages.append({"role": "assistant", "content": response})
-    st.chat_message("assistant").write(response)
-
-# Кнопка очистки истории
-if st.sidebar.button("Очистить историю диалога"):
-    st.session_state.messages = [
-        {"role": "assistant", "content": "История очищена. Чем могу помочь?"}
-    ]
-    st.rerun()
+    st.session_state.messages.append({"role": "assistant", "content": assistant_reply})
+    st.chat_message("assistant").write(assistant_reply)