rag_deepseek_app.py

# -*- coding: utf-8 -*-
"""rag_deepseek_app.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1P-gvSUDPm7r5O1H4H50ljiQIiG1bU_cC
"""

# rag_deepseek_app.py

import streamlit as st
import openai
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer
import PyPDF2
import docx
import io
import re
import os
from typing import List, Dict, Any

class DocumentProcessor:
    def __init__(self):
        self.supported_formats = ['.pdf', '.txt', '.docx']

    def extract_text_from_pdf(self, file_bytes: bytes) -> str:
        try:
            pdf_reader = PyPDF2.PdfReader(io.BytesIO(file_bytes))
            text = ""
            for page in pdf_reader.pages:
                text += page.extract_text() + "\n"
            return text
        except Exception as e:
            st.error(f"Error reading PDF: {str(e)}")
            return ""

    def extract_text_from_docx(self, file_bytes: bytes) -> str:
        try:
            doc = docx.Document(io.BytesIO(file_bytes))
            return "\n".join(paragraph.text for paragraph in doc.paragraphs)
        except Exception as e:
            st.error(f"Error reading DOCX: {str(e)}")
            return ""

    def extract_text_from_txt(self, file_bytes: bytes) -> str:
        try:
            return file_bytes.decode('utf-8')
        except Exception as e:
            st.error(f"Error reading TXT: {str(e)}")
            return ""

    def process_document(self, uploaded_file) -> str:
        file_extension = os.path.splitext(uploaded_file.name)[1].lower()
        file_bytes = uploaded_file.read()

        if file_extension == '.pdf':
            return self.extract_text_from_pdf(file_bytes)
        elif file_extension == '.docx':
            return self.extract_text_from_docx(file_bytes)
        elif file_extension == '.txt':
            return self.extract_text_from_txt(file_bytes)
        else:
            st.error(f"Unsupported file format: {file_extension}")
            return ""

class TextChunker:
    def __init__(self, chunk_size: int = 1000, overlap: int = 200):
        self.chunk_size = chunk_size
        self.overlap = overlap

    def chunk_text(self, text: str) -> List[str]:
        text = re.sub(r'\s+', ' ', text.strip())
        if len(text) <= self.chunk_size:
            return [text]

        chunks = []
        start = 0
        while start < len(text):
            end = start + self.chunk_size
            if end < len(text):
                sentence_end = text.rfind('.', start, end)
                if sentence_end != -1 and sentence_end > start + self.chunk_size // 2:
                    end = sentence_end + 1
                else:
                    word_end = text.rfind(' ', start, end)
                    if word_end != -1 and word_end > start + self.chunk_size // 2:
                        end = word_end
            chunk = text[start:end].strip()
            if chunk:
                chunks.append(chunk)
            start = end - self.overlap
            if start <= 0:
                start = end
        return chunks

class VectorStore:
    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.model = SentenceTransformer(model_name)
        self.index = None
        self.chunks = []
        self.embeddings = None

    def create_embeddings(self, chunks: List[str]) -> np.ndarray:
        with st.spinner("Creating embeddings..."):
            embeddings = self.model.encode(chunks, show_progress_bar=False)
        return embeddings

    def build_index(self, chunks: List[str]):
        self.chunks = chunks
        self.embeddings = self.create_embeddings(chunks)
        dimension = self.embeddings.shape[1]
        self.index = faiss.IndexFlatIP(dimension)
        faiss.normalize_L2(self.embeddings)
        self.index.add(self.embeddings)

    def search(self, query: str, k: int = 3) -> List[Dict[str, Any]]:
        if self.index is None:
            return []
        query_embedding = self.model.encode([query])
        faiss.normalize_L2(query_embedding)
        scores, indices = self.index.search(query_embedding, k)
        results = []
        for i, (score, idx) in enumerate(zip(scores[0], indices[0])):
            if idx < len(self.chunks):
                results.append({'text': self.chunks[idx], 'score': float(score), 'rank': i + 1})
        return results

class DeepSeekRAG:
    def __init__(self, api_key: str):
        self.client = openai.OpenAI(
            api_key=api_key,
            base_url="https://api.deepseek.com"
        )
        self.doc_processor = DocumentProcessor()
        self.chunker = TextChunker()
        self.vector_store = VectorStore()

    def process_document(self, uploaded_file) -> bool:
        try:
            text = self.doc_processor.process_document(uploaded_file)
            if not text:
                return False
            chunks = self.chunker.chunk_text(text)
            if not chunks:
                st.error("No text chunks created from document")
                return False
            self.vector_store.build_index(chunks)
            st.success(f"Document processed successfully! Created {len(chunks)} chunks.")
            return True
        except Exception as e:
            st.error(f"Error processing document: {str(e)}")
            return False

    def answer_question(self, question: str, max_tokens: int = 1000) -> Dict[str, Any]:
        try:
            relevant_chunks = self.vector_store.search(question, k=3)
            if not relevant_chunks:
                return {'answer': "No relevant information found.", 'sources': [], 'error': None}
            context = "\n\n".join([chunk['text'] for chunk in relevant_chunks])
            prompt = f"""Based on the following context from the document, please answer the question. If the answer cannot be found in the context, please say so.

Context:
{context}

Question: {question}

Answer:"""
            response = self.client.chat.completions.create(
                model="deepseek-chat",
                messages=[
                    {"role": "system", "content": "You are a helpful assistant that answers questions based on provided context."},
                    {"role": "user", "content": prompt}
                ],
                max_tokens=max_tokens,
                temperature=0.1
            )
            answer = response.choices[0].message.content
            return {'answer': answer, 'sources': relevant_chunks, 'error': None}
        except Exception as e:
            return {'answer': f"Error generating answer: {str(e)}", 'sources': [], 'error': str(e)}

def main():
    st.set_page_config(page_title="RAG Document Q&A System", page_icon="📚", layout="wide")
    st.title("📚 RAG Document Q&A System")
    st.markdown("Upload a document and ask questions about its content using DeepSeek AI!")

    with st.sidebar:
        st.header("Configuration")
        api_key = st.text_input("DeepSeek API Key", type="password", help="Enter your DeepSeek API key")
        uploaded_file = st.file_uploader("Choose a document", type=['pdf', 'txt', 'docx'])
        max_tokens = st.slider("Max Response Tokens", 100, 2000, 1000)
        chunk_size = st.slider("Chunk Size", 500, 2000, 1000)
        overlap = st.slider("Chunk Overlap", 50, 500, 200)

    if 'rag_system' not in st.session_state:
        st.session_state.rag_system = None
    if 'document_processed' not in st.session_state:
        st.session_state.document_processed = False
    if 'chat_history' not in st.session_state:
        st.session_state.chat_history = []

    col1, col2 = st.columns([2, 1])

    with col1:
        if not api_key:
            st.warning("Please enter your DeepSeek API key in the sidebar.")
        elif not uploaded_file:
            st.info("Please upload a document.")
        else:
            if st.session_state.rag_system is None:
                st.session_state.rag_system = DeepSeekRAG(api_key)
                st.session_state.rag_system.chunker = TextChunker(chunk_size, overlap)
            if not st.session_state.document_processed:
                with st.spinner("Processing document..."):
                    success = st.session_state.rag_system.process_document(uploaded_file)
                    st.session_state.document_processed = success
            if st.session_state.document_processed:
                st.success("Document is ready for questions!")
                question = st.text_input("Ask a question about the document:", placeholder="What is the main topic?")
                if st.button("Get Answer", disabled=not question):
                    with st.spinner("Generating answer..."):
                        result = st.session_state.rag_system.answer_question(question, max_tokens)
                        st.session_state.chat_history.append({
                            'question': question,
                            'answer': result['answer'],
                            'sources': result['sources']
                        })
                if st.session_state.chat_history:
                    st.markdown("### Chat History")
                    for i, chat in enumerate(reversed(st.session_state.chat_history)):
                        with st.expander(f"Q: {chat['question'][:50]}...", expanded=(i == 0)):
                            st.markdown(f"**Question:** {chat['question']}")
                            st.markdown(f"**Answer:** {chat['answer']}")
                            if chat['sources']:
                                st.markdown("**Sources:**")
                                for j, source in enumerate(chat['sources']):
                                    st.markdown(f"**Chunk {j+1}** (Score: {source['score']:.3f})")
                                    st.text(source['text'][:200] + "..." if len(source['text']) > 200 else source['text'])

    with col2:
        st.markdown("### About")
        st.markdown("""
        This RAG (Retrieval-Augmented Generation) system:
        1. Processes your document into chunks
        2. Creates vector embeddings for semantic search
        3. Retrieves relevant chunks
        4. Generates answers using DeepSeek AI
        """)
        if st.session_state.document_processed:
            st.metric("Chunks Created", len(st.session_state.rag_system.vector_store.chunks))
            st.metric("Questions Asked", len(st.session_state.chat_history))
        if st.button("Clear History"):
            st.session_state.chat_history = []
            st.session_state.document_processed = False
            st.session_state.rag_system = None
            st.rerun()

if __name__ == "__main__":
    main()