Create app.py

app.py

@@ -0,0 +1,363 @@
+import os
+import streamlit as st
+import numpy as np
+import pandas as pd
+from sentence_transformers import SentenceTransformer
+from groq import Groq
+import faiss
+import pickle
+from typing import List, Dict, Tuple
+import PyPDF2
+import docx
+from io import BytesIO
+import time
+
+# Initialize Groq client
+def init_groq_client(api_key: str):
+    """Initialize Groq client with API key"""
+    return Groq(api_key=api_key)
+
+# Initialize embedding model
+@st.cache_resource
+def load_embedding_model():
+    """Load and cache the sentence transformer model"""
+    return SentenceTransformer('all-MiniLM-L6-v2')
+
+# Document processing functions
+def extract_text_from_pdf(file):
+    """Extract text from PDF file"""
+    pdf_reader = PyPDF2.PdfReader(file)
+    text = ""
+    for page in pdf_reader.pages:
+        text += page.extract_text()
+    return text
+
+def extract_text_from_docx(file):
+    """Extract text from DOCX file"""
+    doc = docx.Document(file)
+    text = ""
+    for paragraph in doc.paragraphs:
+        text += paragraph.text + "\n"
+    return text
+
+def extract_text_from_txt(file):
+    """Extract text from TXT file"""
+    return str(file.read(), "utf-8")
+
+def chunk_text(text: str, chunk_size: int = 500, overlap: int = 50) -> List[str]:
+    """Split text into overlapping chunks"""
+    words = text.split()
+    chunks = []
+    
+    for i in range(0, len(words), chunk_size - overlap):
+        chunk = ' '.join(words[i:i + chunk_size])
+        chunks.append(chunk)
+        
+        if i + chunk_size >= len(words):
+            break
+    
+    return chunks
+
+# Vector store class
+class VectorStore:
+    def __init__(self, embedding_model):
+        self.embedding_model = embedding_model
+        self.documents = []
+        self.embeddings = []
+        self.index = None
+        
+    def add_documents(self, documents: List[str]):
+        """Add documents to the vector store"""
+        self.documents.extend(documents)
+        
+        # Generate embeddings
+        new_embeddings = self.embedding_model.encode(documents)
+        
+        if len(self.embeddings) == 0:
+            self.embeddings = new_embeddings
+        else:
+            self.embeddings = np.vstack([self.embeddings, new_embeddings])
+        
+        # Build/update FAISS index
+        self._build_index()
+    
+    def _build_index(self):
+        """Build FAISS index for similarity search"""
+        if len(self.embeddings) > 0:
+            dimension = self.embeddings.shape[1]
+            self.index = faiss.IndexFlatIP(dimension)  # Inner product for similarity
+            
+            # Normalize embeddings for cosine similarity
+            normalized_embeddings = self.embeddings / np.linalg.norm(
+                self.embeddings, axis=1, keepdims=True
+            )
+            self.index.add(normalized_embeddings.astype('float32'))
+    
+    def search(self, query: str, top_k: int = 3) -> List[Tuple[str, float]]:
+        """Search for similar documents"""
+        if self.index is None or len(self.documents) == 0:
+            return []
+        
+        # Encode query
+        query_embedding = self.embedding_model.encode([query])
+        query_embedding = query_embedding / np.linalg.norm(query_embedding, axis=1, keepdims=True)
+        
+        # Search
+        scores, indices = self.index.search(query_embedding.astype('float32'), top_k)
+        
+        results = []
+        for score, idx in zip(scores[0], indices[0]):
+            if idx < len(self.documents):
+                results.append((self.documents[idx], float(score)))
+        
+        return results
+    
+    def save(self, filepath: str):
+        """Save vector store to file"""
+        data = {
+            'documents': self.documents,
+            'embeddings': self.embeddings.tolist() if len(self.embeddings) > 0 else []
+        }
+        with open(filepath, 'wb') as f:
+            pickle.dump(data, f)
+    
+    def load(self, filepath: str):
+        """Load vector store from file"""
+        with open(filepath, 'rb') as f:
+            data = pickle.load(f)
+        
+        self.documents = data['documents']
+        if data['embeddings']:
+            self.embeddings = np.array(data['embeddings'])
+            self._build_index()
+
+# RAG class
+class RAGSystem:
+    def __init__(self, groq_client, embedding_model):
+        self.groq_client = groq_client
+        self.vector_store = VectorStore(embedding_model)
+        
+    def add_documents(self, documents: List[str]):
+        """Add documents to the knowledge base"""
+        self.vector_store.add_documents(documents)
+    
+    def query(self, question: str, model: str = "llama-3.3-70b-versatile", top_k: int = 3) -> Dict:
+        """Answer a question using RAG"""
+        # Retrieve relevant documents
+        retrieved_docs = self.vector_store.search(question, top_k=top_k)
+        
+        if not retrieved_docs:
+            return {
+                "answer": "I don't have any relevant information to answer your question.",
+                "sources": [],
+                "confidence": 0.0
+            }
+        
+        # Prepare context
+        context = "\n\n".join([doc for doc, score in retrieved_docs])
+        
+        # Create prompt
+        prompt = f"""Based on the following context, answer the question. If the answer is not in the context, say "I don't have enough information to answer this question."
+
+Context:
+{context}
+
+Question: {question}
+
+Answer:"""
+        
+        try:
+            # Get response from Groq
+            chat_completion = self.groq_client.chat.completions.create(
+                messages=[
+                    {
+                        "role": "user",
+                        "content": prompt,
+                    }
+                ],
+                model=model,
+                temperature=0.1,
+                max_tokens=1000,
+            )
+            
+            answer = chat_completion.choices[0].message.content
+            
+            return {
+                "answer": answer,
+                "sources": [{"text": doc[:200] + "...", "score": score} 
+                           for doc, score in retrieved_docs],
+                "confidence": max([score for _, score in retrieved_docs]) if retrieved_docs else 0.0
+            }
+            
+        except Exception as e:
+            return {
+                "answer": f"Error generating response: {str(e)}",
+                "sources": [],
+                "confidence": 0.0
+            }
+
+# Streamlit App
+def main():
+    st.set_page_config(
+        page_title="RAG App with Groq",
+        page_icon="🤖",
+        layout="wide",
+        initial_sidebar_state="expanded"
+    )
+    
+    st.title("🤖 RAG App with Groq & Sentence Transformers")
+    st.markdown("Ask questions about your documents using open-source models!")
+    
+    # Sidebar
+    st.sidebar.header("⚙️ Configuration")
+    
+    # API Key input
+    api_key = st.sidebar.text_input(
+        "Groq API Key",
+        value=os.getenv("GROQ_API_KEY", ""),
+        type="password",
+        help="Enter your Groq API key"
+    )
+    
+    # Model selection
+    model_options = [
+        "llama-3.3-70b-versatile",
+        "llama-3.1-70b-versatile",
+        "llama-3.1-8b-instant",
+        "mixtral-8x7b-32768"
+    ]
+    selected_model = st.sidebar.selectbox("Select Model", model_options)
+    
+    # Number of retrieved documents
+    top_k = st.sidebar.slider("Number of retrieved documents", 1, 10, 3)
+    
+    # Initialize components
+    if api_key:
+        try:
+            groq_client = init_groq_client(api_key)
+            embedding_model = load_embedding_model()
+            
+            # Initialize session state
+            if 'rag_system' not in st.session_state:
+                st.session_state.rag_system = RAGSystem(groq_client, embedding_model)
+            
+            # Main content area
+            col1, col2 = st.columns([1, 1])
+            
+            with col1:
+                st.header("📁 Document Upload")
+                
+                uploaded_files = st.file_uploader(
+                    "Upload your documents",
+                    type=['pdf', 'docx', 'txt'],
+                    accept_multiple_files=True,
+                    help="Supported formats: PDF, DOCX, TXT"
+                )
+                
+                if uploaded_files:
+                    if st.button("Process Documents", type="primary"):
+                        with st.spinner("Processing documents..."):
+                            all_chunks = []
+                            
+                            for file in uploaded_files:
+                                # Extract text based on file type
+                                if file.type == "application/pdf":
+                                    text = extract_text_from_pdf(file)
+                                elif file.type == "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
+                                    text = extract_text_from_docx(file)
+                                elif file.type == "text/plain":
+                                    text = extract_text_from_txt(file)
+                                else:
+                                    st.error(f"Unsupported file type: {file.type}")
+                                    continue
+                                
+                                # Chunk the text
+                                chunks = chunk_text(text, chunk_size=500, overlap=50)
+                                all_chunks.extend(chunks)
+                                
+                                st.success(f"✅ Processed {file.name}: {len(chunks)} chunks")
+                            
+                            # Add to RAG system
+                            if all_chunks:
+                                st.session_state.rag_system.add_documents(all_chunks)
+                                st.success(f"🎉 Added {len(all_chunks)} chunks to knowledge base!")
+                
+                # Display document stats
+                if hasattr(st.session_state.rag_system, 'vector_store') and len(st.session_state.rag_system.vector_store.documents) > 0:
+                    st.info(f"📊 Knowledge Base: {len(st.session_state.rag_system.vector_store.documents)} chunks")
+            
+            with col2:
+                st.header("💬 Ask Questions")
+                
+                # Chat interface
+                if "messages" not in st.session_state:
+                    st.session_state.messages = []
+                
+                # Display chat history
+                for message in st.session_state.messages:
+                    with st.chat_message(message["role"]):
+                        st.write(message["content"])
+                        if message["role"] == "assistant" and "sources" in message:
+                            with st.expander("📚 Sources"):
+                                for i, source in enumerate(message["sources"]):
+                                    st.write(f"**Source {i+1}** (Score: {source['score']:.3f})")
+                                    st.write(source["text"])
+                
+                # Chat input
+                if prompt := st.chat_input("Ask a question about your documents..."):
+                    # Add user message
+                    st.session_state.messages.append({"role": "user", "content": prompt})
+                    
+                    with st.chat_message("user"):
+                        st.write(prompt)
+                    
+                    # Generate response
+                    with st.chat_message("assistant"):
+                        with st.spinner("Thinking..."):
+                            response = st.session_state.rag_system.query(
+                                prompt, 
+                                model=selected_model, 
+                                top_k=top_k
+                            )
+                        
+                        st.write(response["answer"])
+                        
+                        # Show sources
+                        if response["sources"]:
+                            with st.expander("📚 Sources"):
+                                for i, source in enumerate(response["sources"]):
+                                    st.write(f"**Source {i+1}** (Score: {source['score']:.3f})")
+                                    st.write(source["text"])
+                        
+                        # Add to chat history
+                        st.session_state.messages.append({
+                            "role": "assistant", 
+                            "content": response["answer"],
+                            "sources": response["sources"]
+                        })
+                
+                # Clear chat button
+                if st.button("🗑️ Clear Chat"):
+                    st.session_state.messages = []
+                    st.rerun()
+        
+        except Exception as e:
+            st.error(f"Error initializing components: {str(e)}")
+    
+    else:
+        st.warning("Please enter your Groq API key in the sidebar to get started.")
+    
+    # Footer
+    st.sidebar.markdown("---")
+    st.sidebar.markdown(
+        """
+        **About this app:**
+        - Uses Groq for fast inference
+        - Sentence Transformers for embeddings
+        - FAISS for vector search
+        - Supports PDF, DOCX, TXT files
+        """
+    )
+
+if __name__ == "__main__":
+    main()