Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,169 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 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
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+# Конфигурация
+DATA_DIR = "data"
+INDEX_DIR = "faiss_index"
+MODEL_NAME = "IlyaGusev/saiga_llama3_8b"
+
+# Инициализация модели
+@st.cache_resource
+def load_llm():
+    return pipeline(
+        "text-generation",
+        model=MODEL_NAME,
+        device_map="auto",
+        model_kwargs={"torch_dtype": "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("База знаний обновлена!")
+
+# История диалога
+if "messages" not in st.session_state:
+    st.session_state.messages = [
+        {"role": "assistant", "content": "Привет! Задайте вопрос о магистратуре, и я отвечу на основе документов."}
+    ]
+
+# Отображение истории
+for msg in st.session_state.messages:
+    st.chat_message(msg["role"]).write(msg["content"])
+
+# Обработка ввода
+if prompt := st.chat_input("Ваш вопрос о магистратуре..."):
+    # Добавляем вопрос в историю
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    st.chat_message("user").write(prompt)
+    
+    # Генерация ответа с RAG
+    with st.spinner("Ищу информацию..."):
+        try:
+            response = generate_with_rag(
+                prompt, 
+                st.session_state.messages[-5:]  # Последние 5 сообщений как контекст
+            )
+        except Exception as e:
+            response = 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()