Create app.py

app.py

@@ -0,0 +1,444 @@
+import streamlit as st
+import os
+import tempfile
+import pickle
+from typing import List, Dict, Any
+import numpy as np
+from pathlib import Path
+
+# Document processing
+import PyPDF2
+import docx
+from sentence_transformers import SentenceTransformer
+import faiss
+
+# Groq API
+from groq import Groq
+
+# Text processing
+import nltk
+from nltk.tokenize import sent_tokenize
+import re
+
+# Download required NLTK data
+try:
+    nltk.data.find('tokenizers/punkt')
+except LookupError:
+    nltk.download('punkt')
+
+class DocumentProcessor:
+    """Handles document upload and text extraction"""
+    
+    @staticmethod
+    def extract_text_from_pdf(file_path: str) -> str:
+        """Extract text from PDF file"""
+        text = ""
+        try:
+            with open(file_path, 'rb') as file:
+                pdf_reader = PyPDF2.PdfReader(file)
+                for page in pdf_reader.pages:
+                    text += page.extract_text() + "\n"
+        except Exception as e:
+            st.error(f"Error reading PDF: {str(e)}")
+        return text
+    
+    @staticmethod
+    def extract_text_from_docx(file_path: str) -> str:
+        """Extract text from DOCX file"""
+        text = ""
+        try:
+            doc = docx.Document(file_path)
+            for paragraph in doc.paragraphs:
+                text += paragraph.text + "\n"
+        except Exception as e:
+            st.error(f"Error reading DOCX: {str(e)}")
+        return text
+    
+    @staticmethod
+    def extract_text_from_txt(file_path: str) -> str:
+        """Extract text from TXT file"""
+        text = ""
+        try:
+            with open(file_path, 'r', encoding='utf-8') as file:
+                text = file.read()
+        except Exception as e:
+            st.error(f"Error reading TXT: {str(e)}")
+        return text
+    
+    def process_uploaded_file(self, uploaded_file) -> str:
+        """Process uploaded file and extract text"""
+        if uploaded_file is None:
+            return ""
+        
+        # Save uploaded file temporarily
+        with tempfile.NamedTemporaryFile(delete=False, suffix=f".{uploaded_file.name.split('.')[-1]}") as tmp_file:
+            tmp_file.write(uploaded_file.getvalue())
+            tmp_file_path = tmp_file.name
+        
+        try:
+            file_extension = uploaded_file.name.split('.')[-1].lower()
+            
+            if file_extension == 'pdf':
+                text = self.extract_text_from_pdf(tmp_file_path)
+            elif file_extension == 'docx':
+                text = self.extract_text_from_docx(tmp_file_path)
+            elif file_extension == 'txt':
+                text = self.extract_text_from_txt(tmp_file_path)
+            else:
+                st.error(f"Unsupported file type: {file_extension}")
+                return ""
+            
+            return text
+        finally:
+            # Clean up temporary file
+            os.unlink(tmp_file_path)
+
+class TextChunker:
+    """Handles text chunking and preprocessing"""
+    
+    def __init__(self, chunk_size: int = 1000, overlap: int = 200):
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+    
+    def clean_text(self, text: str) -> str:
+        """Clean and preprocess text"""
+        # Remove extra whitespace
+        text = re.sub(r'\s+', ' ', text)
+        # Remove special characters but keep punctuation
+        text = re.sub(r'[^\w\s\.\!\?\,\;\:\-\(\)]', '', text)
+        return text.strip()
+    
+    def create_chunks(self, text: str) -> List[str]:
+        """Create overlapping chunks from text"""
+        cleaned_text = self.clean_text(text)
+        
+        # Split into sentences first
+        sentences = sent_tokenize(cleaned_text)
+        
+        chunks = []
+        current_chunk = ""
+        
+        for sentence in sentences:
+            # If adding this sentence would exceed chunk size, start a new chunk
+            if len(current_chunk) + len(sentence) > self.chunk_size:
+                if current_chunk:
+                    chunks.append(current_chunk.strip())
+                    
+                    # Create overlap by keeping last part of current chunk
+                    words = current_chunk.split()
+                    if len(words) > 20:  # Keep last 20 words for overlap
+                        current_chunk = " ".join(words[-20:]) + " " + sentence
+                    else:
+                        current_chunk = sentence
+                else:
+                    current_chunk = sentence
+            else:
+                current_chunk += " " + sentence
+        
+        # Add the last chunk
+        if current_chunk:
+            chunks.append(current_chunk.strip())
+        
+        return chunks
+
+class VectorDatabase:
+    """Handles vector embeddings and FAISS operations"""
+    
+    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
+        self.embedding_model = SentenceTransformer(model_name)
+        self.dimension = self.embedding_model.get_sentence_embedding_dimension()
+        self.index = faiss.IndexFlatIP(self.dimension)  # Inner product for similarity
+        self.chunks = []
+        self.embeddings = None
+    
+    def create_embeddings(self, chunks: List[str]) -> np.ndarray:
+        """Create embeddings for text chunks"""
+        with st.spinner("Creating embeddings..."):
+            embeddings = self.embedding_model.encode(chunks, show_progress_bar=True)
+            # Normalize embeddings for cosine similarity
+            embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
+        return embeddings
+    
+    def add_documents(self, chunks: List[str]):
+        """Add documents to the vector database"""
+        if not chunks:
+            return
+        
+        self.chunks.extend(chunks)
+        embeddings = self.create_embeddings(chunks)
+        
+        if self.embeddings is None:
+            self.embeddings = embeddings
+        else:
+            self.embeddings = np.vstack([self.embeddings, embeddings])
+        
+        # Add to FAISS index
+        self.index.add(embeddings.astype(np.float32))
+        
+        st.success(f"Added {len(chunks)} chunks to vector database")
+    
+    def search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
+        """Search for similar documents"""
+        if self.index.ntotal == 0:
+            return []
+        
+        # Create query embedding
+        query_embedding = self.embedding_model.encode([query])
+        query_embedding = query_embedding / np.linalg.norm(query_embedding)
+        
+        # Search in FAISS
+        scores, indices = self.index.search(query_embedding.astype(np.float32), k)
+        
+        results = []
+        for i, (score, idx) in enumerate(zip(scores[0], indices[0])):
+            if idx < len(self.chunks):
+                results.append({
+                    'chunk': self.chunks[idx],
+                    'score': float(score),
+                    'rank': i + 1
+                })
+        
+        return results
+    
+    def save_database(self, filepath: str):
+        """Save the vector database to disk"""
+        data = {
+            'chunks': self.chunks,
+            'embeddings': self.embeddings,
+            'index': faiss.serialize_index(self.index)
+        }
+        
+        with open(filepath, 'wb') as f:
+            pickle.dump(data, f)
+    
+    def load_database(self, filepath: str):
+        """Load the vector database from disk"""
+        try:
+            with open(filepath, 'rb') as f:
+                data = pickle.load(f)
+            
+            self.chunks = data['chunks']
+            self.embeddings = data['embeddings']
+            self.index = faiss.deserialize_index(data['index'])
+            
+            return True
+        except Exception as e:
+            st.error(f"Error loading database: {str(e)}")
+            return False
+
+class RAGSystem:
+    """Main RAG system that combines retrieval and generation"""
+    
+    def __init__(self, groq_api_key: str):
+        self.groq_client = Groq(api_key=groq_api_key)
+        self.vector_db = VectorDatabase()
+        self.doc_processor = DocumentProcessor()
+        self.text_chunker = TextChunker()
+    
+    def process_document(self, uploaded_file):
+        """Process uploaded document and add to vector database"""
+        # Extract text from document
+        text = self.doc_processor.process_uploaded_file(uploaded_file)
+        
+        if not text:
+            st.error("No text extracted from document")
+            return False
+        
+        # Create chunks
+        chunks = self.text_chunker.create_chunks(text)
+        
+        if not chunks:
+            st.error("No chunks created from text")
+            return False
+        
+        # Add to vector database
+        self.vector_db.add_documents(chunks)
+        
+        return True
+    
+    def generate_response(self, query: str, context: str, model: str = "llama-3.3-70b-versatile") -> str:
+        """Generate response using Groq API"""
+        
+        prompt = f"""
+        Based on the following context, please answer the question. If the answer is not in the context, say "I don't have enough information to answer this question based on the provided documents."
+
+        Context:
+        {context}
+
+        Question: {query}
+
+        Answer:
+        """
+        
+        try:
+            chat_completion = self.groq_client.chat.completions.create(
+                messages=[
+                    {
+                        "role": "system",
+                        "content": "You are a helpful assistant that answers questions based on provided context. Be accurate and concise."
+                    },
+                    {
+                        "role": "user",
+                        "content": prompt
+                    }
+                ],
+                model=model,
+                temperature=0.1,
+                max_tokens=1000
+            )
+            
+            return chat_completion.choices[0].message.content
+        except Exception as e:
+            return f"Error generating response: {str(e)}"
+    
+    def query(self, question: str, model: str = "llama-3.3-70b-versatile") -> Dict[str, Any]:
+        """Query the RAG system"""
+        # Retrieve relevant documents
+        search_results = self.vector_db.search(question, k=3)
+        
+        if not search_results:
+            return {
+                'answer': "No relevant documents found. Please upload some documents first.",
+                'sources': []
+            }
+        
+        # Combine contexts
+        context = "\n\n".join([result['chunk'] for result in search_results])
+        
+        # Generate response
+        answer = self.generate_response(question, context, model)
+        
+        return {
+            'answer': answer,
+            'sources': search_results
+        }
+
+def main():
+    st.set_page_config(
+        page_title="RAG Application",
+        page_icon="🔍",
+        layout="wide",
+        initial_sidebar_state="expanded"
+    )
+    
+    st.title("🔍 RAG Application")
+    st.markdown("**Upload documents and ask questions using AI-powered search and generation**")
+    
+    # Initialize session state
+    if 'rag_system' not in st.session_state:
+        st.session_state.rag_system = None
+    if 'documents_processed' not in st.session_state:
+        st.session_state.documents_processed = 0
+    
+    # Sidebar for configuration
+    with st.sidebar:
+        st.header("⚙️ Configuration")
+        
+        # API Key input
+        groq_api_key = st.text_input(
+            "Groq API Key",
+            type="password",
+            help="Enter your Groq API key"
+        )
+        
+        if not groq_api_key:
+            st.warning("Please enter your Groq API key to continue")
+            st.stop()
+        
+        # Model selection
+        model_options = [
+            "llama-3.3-70b-versatile",
+            "llama-3.2-90b-text-preview",
+            "llama-3.1-70b-versatile",
+            "mixtral-8x7b-32768",
+            "gemma2-9b-it"
+        ]
+        
+        selected_model = st.selectbox(
+            "Select Model",
+            model_options,
+            index=0
+        )
+        
+        # Initialize RAG system
+        if st.session_state.rag_system is None:
+            try:
+                st.session_state.rag_system = RAGSystem(groq_api_key)
+                st.success("RAG system initialized!")
+            except Exception as e:
+                st.error(f"Error initializing RAG system: {str(e)}")
+                st.stop()
+        
+        st.header("📊 Statistics")
+        st.metric("Documents Processed", st.session_state.documents_processed)
+        st.metric("Chunks in Database", len(st.session_state.rag_system.vector_db.chunks))
+    
+    # Main content area
+    col1, col2 = st.columns([1, 2])
+    
+    with col1:
+        st.header("📄 Document Upload")
+        
+        uploaded_files = st.file_uploader(
+            "Upload documents",
+            accept_multiple_files=True,
+            type=['pdf', 'docx', 'txt'],
+            help="Upload PDF, DOCX, or TXT files"
+        )
+        
+        if uploaded_files:
+            for uploaded_file in uploaded_files:
+                if st.button(f"Process {uploaded_file.name}"):
+                    with st.spinner(f"Processing {uploaded_file.name}..."):
+                        success = st.session_state.rag_system.process_document(uploaded_file)
+                        if success:
+                            st.session_state.documents_processed += 1
+                            st.success(f"Successfully processed {uploaded_file.name}")
+                        else:
+                            st.error(f"Failed to process {uploaded_file.name}")
+    
+    with col2:
+        st.header("💬 Ask Questions")
+        
+        if len(st.session_state.rag_system.vector_db.chunks) == 0:
+            st.info("Please upload and process documents before asking questions.")
+        else:
+            question = st.text_input(
+                "Enter your question:",
+                placeholder="What is this document about?"
+            )
+            
+            if st.button("Ask Question") and question:
+                with st.spinner("Generating answer..."):
+                    response = st.session_state.rag_system.query(question, selected_model)
+                    
+                    st.subheader("Answer:")
+                    st.write(response['answer'])
+                    
+                    if response['sources']:
+                        st.subheader("Sources:")
+                        for i, source in enumerate(response['sources']):
+                            with st.expander(f"Source {i+1} (Score: {source['score']:.3f})"):
+                                st.write(source['chunk'])
+    
+    # Additional features
+    st.header("🔧 Additional Features")
+    
+    col3, col4 = st.columns(2)
+    
+    with col3:
+        if st.button("Clear Database"):
+            st.session_state.rag_system.vector_db = VectorDatabase()
+            st.session_state.documents_processed = 0
+            st.success("Database cleared successfully!")
+    
+    with col4:
+        if st.button("Save Database"):
+            if len(st.session_state.rag_system.vector_db.chunks) > 0:
+                st.session_state.rag_system.vector_db.save_database("rag_database.pkl")
+                st.success("Database saved successfully!")
+            else:
+                st.warning("No data to save")
+
+if __name__ == "__main__":
+    main()