Create RAGSample.py

src/RAGSample.py

@@ -0,0 +1,388 @@
+import os
+from dotenv import load_dotenv
+from langchain_community.document_loaders import WebBaseLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import Chroma
+from langchain_ollama import ChatOllama
+from langchain.prompts import PromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.retrievers import BaseRetriever
+from langchain_core.runnables import Runnable
+from langchain_core.documents import Document
+from langchain_core.embeddings import Embeddings
+import chromadb
+import numpy as np
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+import pandas as pd
+from typing import Optional, List
+import re
+
+# Disable ChromaDB telemetry to avoid the error
+os.environ["ANONYMIZED_TELEMETRY"] = "False"
+os.environ["CHROMA_SERVER_HOST"] = "localhost"
+os.environ["CHROMA_SERVER_HTTP_PORT"] = "8000"
+
+
+class ImprovedTFIDFEmbeddings(Embeddings):
+    """Improved TF-IDF based embedding function with better preprocessing."""
+    
+    def __init__(self):
+        self.vectorizer = TfidfVectorizer(
+            max_features=5000,
+            stop_words='english',
+            ngram_range=(1, 3),
+            min_df=1,
+            max_df=0.85,
+            lowercase=True,
+            strip_accents='unicode',
+            analyzer='word'
+        )
+        self.fitted = False
+        self.documents = []
+    
+    def embed_documents(self, texts):
+        """Create embeddings for a list of texts."""
+        if not self.fitted:
+            self.documents = texts
+            self.vectorizer.fit(texts)
+            self.fitted = True
+        
+        # Transform texts to TF-IDF vectors
+        tfidf_matrix = self.vectorizer.transform(texts)
+        
+        # Convert to dense arrays and normalize
+        embeddings = []
+        for i in range(tfidf_matrix.shape[0]):
+            embedding = tfidf_matrix[i].toarray().flatten()
+            # Normalize the embedding
+            norm = np.linalg.norm(embedding)
+            if norm > 0:
+                embedding = embedding / norm
+            # Pad or truncate to 512 dimensions
+            if len(embedding) < 512:
+                embedding = np.pad(embedding, (0, 512 - len(embedding)))
+            else:
+                embedding = embedding[:512]
+            embeddings.append(embedding.tolist())
+        
+        return embeddings
+    
+    def embed_query(self, text):
+        """Create embedding for a single query text."""
+        if not self.fitted:
+            # If not fitted, fit with just this text
+            self.vectorizer.fit([text])
+            self.fitted = True
+        
+        # Transform query to TF-IDF vector
+        tfidf_matrix = self.vectorizer.transform([text])
+        embedding = tfidf_matrix[0].toarray().flatten()
+        
+        # Normalize the embedding
+        norm = np.linalg.norm(embedding)
+        if norm > 0:
+            embedding = embedding / norm
+        
+        # Pad or truncate to 512 dimensions
+        if len(embedding) < 512:
+            embedding = np.pad(embedding, (0, 512 - len(embedding)))
+        else:
+            embedding = embedding[:512]
+        
+        return embedding.tolist()
+
+
+class SmartFAQRetriever(BaseRetriever):
+    """Smart retriever optimized for FAQ datasets with semantic similarity."""
+    
+    def __init__(self, documents: List[Document], k: int = 4):
+        super().__init__()
+        self._documents = documents
+        self._k = k
+        self._vectorizer = None  # Use private attribute
+    
+    @property
+    def documents(self):
+        return self._documents
+    
+    @property
+    def k(self):
+        return self._k
+    
+    def _get_relevant_documents(self, query: str) -> List[Document]:
+        """Retrieve documents based on semantic similarity."""
+        # Ensure vectorizer is fitted
+        if not hasattr(self, '_vectorizer') or self._vectorizer is None or not hasattr(self._vectorizer, 'vocabulary_') or not self._vectorizer.vocabulary_:
+            print("[SmartFAQRetriever] Fitting vectorizer...")
+            self._vectorizer = TfidfVectorizer(
+                max_features=3000,
+                stop_words='english',
+                ngram_range=(1, 2),
+                min_df=1,
+                max_df=0.9
+            )
+            questions = []
+            for doc in self._documents:
+                if "QUESTION:" in doc.page_content:
+                    question_part = doc.page_content.split("ANSWER:")[0]
+                    question = question_part.replace("QUESTION:", "").strip()
+                    questions.append(question)
+                else:
+                    questions.append(doc.page_content)
+            self._vectorizer.fit(questions)
+        query_lower = query.lower().strip()
+        
+        # Extract questions from documents
+        questions = []
+        for doc in self._documents:
+            if "QUESTION:" in doc.page_content:
+                question_part = doc.page_content.split("ANSWER:")[0]
+                question = question_part.replace("QUESTION:", "").strip()
+                questions.append(question)
+            else:
+                questions.append(doc.page_content)
+        
+        # Transform query and questions to TF-IDF vectors
+        query_vector = self._vectorizer.transform([query_lower])
+        question_vectors = self._vectorizer.transform(questions)
+        
+        # Calculate cosine similarities
+        similarities = cosine_similarity(query_vector, question_vectors).flatten()
+        
+        # Get top k documents
+        top_indices = similarities.argsort()[-self._k:][::-1]
+        
+        # Return documents with highest similarity scores
+        relevant_docs = [self._documents[i] for i in top_indices if similarities[i] > 0.1]
+        
+        if not relevant_docs:
+            # Fallback to first k documents if no good matches
+            relevant_docs = self._documents[:self._k]
+        
+        return relevant_docs
+    
+    async def _aget_relevant_documents(self, query: str) -> List[Document]:
+        """Async version of get_relevant_documents."""
+        return self._get_relevant_documents(query)
+
+def setup_retriever(use_kaggle_data: bool = False, kaggle_dataset: Optional[str] = None, 
+                   kaggle_username: Optional[str] = None, kaggle_key: Optional[str] = None,
+                   use_local_mental_health_data: bool = False) -> BaseRetriever:
+    """
+    Creates a vector store with documents from test data, Kaggle datasets, or local mental health data.
+    
+    Args:
+        use_kaggle_data: Whether to load Kaggle data instead of test documents
+        kaggle_dataset: Kaggle dataset name (e.g., 'username/dataset-name')
+        kaggle_username: Your Kaggle username (optional if using kaggle.json)
+        kaggle_key: Your Kaggle API key (optional if using kaggle.json)
+        use_local_mental_health_data: Whether to load local mental health FAQ data
+    """
+    print("Setting up the retriever...")
+    
+    if use_local_mental_health_data:
+        try:
+            print("Loading mental health FAQ data from local file...")
+            mental_health_file = "data/Mental_Health_FAQ.csv"
+            
+            if not os.path.exists(mental_health_file):
+                print(f"Mental health FAQ file not found: {mental_health_file}")
+                use_local_mental_health_data = False
+            else:
+                # Load mental health FAQ data
+                df = pd.read_csv(mental_health_file)
+                documents = []
+                
+                for _, row in df.iterrows():
+                    question = row['Questions']
+                    answer = row['Answers']
+                    # Create document in FAQ format
+                    content = f"QUESTION: {question}\nANSWER: {answer}"
+                    documents.append(Document(page_content=content))
+                
+                print(f"Loaded {len(documents)} mental health FAQ documents")
+                for i, doc in enumerate(documents[:3]):
+                    print(f"Sample FAQ {i+1}: {doc.page_content[:200]}...")
+                
+        except Exception as e:
+            print(f"Error loading mental health data: {e}")
+            use_local_mental_health_data = False
+    
+    if use_kaggle_data and kaggle_dataset:
+        try:
+            from src.kaggle_loader import KaggleDataLoader
+            
+            print(f"Loading Kaggle dataset: {kaggle_dataset}")
+            # Create loader without parameters - it will auto-load from kaggle.json
+            loader = KaggleDataLoader()
+            
+            # Download the dataset
+            dataset_path = loader.download_dataset(kaggle_dataset)
+            
+            # Load documents based on file type - only process files from this specific dataset
+            documents = []
+            
+            # Get the dataset name to identify the correct files
+            dataset_name = kaggle_dataset.split('/')[-1]
+            print(f"Processing files in dataset directory: {dataset_path}")
+            
+            for file in os.listdir(dataset_path):
+                file_path = os.path.join(dataset_path, file)
+                
+                if file.endswith('.csv'):
+                    print(f"Loading CSV file: {file}")
+                    # For FAQ datasets, use the improved loading method
+                    if 'faq' in file.lower() or 'mental' in file.lower():
+                        documents.extend(loader.load_csv_dataset(file_path, [], chunk_size=50))
+                    else:
+                        # For other CSV files, use first few columns as text
+                        df = pd.read_csv(file_path)
+                        text_columns = df.columns[:3].tolist()  # Use first 3 columns
+                        documents.extend(loader.load_csv_dataset(file_path, text_columns, chunk_size=50))
+                
+                elif file.endswith('.json'):
+                    print(f"Loading JSON file: {file}")
+                    documents.extend(loader.load_json_dataset(file_path))
+                
+                elif file.endswith('.txt'):
+                    print(f"Loading text file: {file}")
+                    documents.extend(loader.load_text_dataset(file_path))
+            
+            print(f"Loaded {len(documents)} documents from Kaggle dataset")
+            for i, doc in enumerate(documents[:3]):
+                print(f"Sample doc {i+1}: {doc.page_content[:200]}")
+            
+        except Exception as e:
+            print(f"Error loading Kaggle data: {e}")
+            print("Falling back to test documents...")
+            use_kaggle_data = False
+    
+    if not use_kaggle_data and not use_local_mental_health_data:
+        # No test documents - use mental health data as default
+        print("No specific data source specified, loading mental health FAQ data as default...")
+        try:
+            mental_health_file = "data/Mental_Health_FAQ.csv"
+            
+            if not os.path.exists(mental_health_file):
+                raise FileNotFoundError(f"Mental health FAQ file not found: {mental_health_file}")
+            
+            # Load mental health FAQ data
+            df = pd.read_csv(mental_health_file)
+            documents = []
+            
+            for _, row in df.iterrows():
+                question = row['Questions']
+                answer = row['Answers']
+                # Create document in FAQ format
+                content = f"QUESTION: {question}\nANSWER: {answer}"
+                documents.append(Document(page_content=content))
+            
+            print(f"Loaded {len(documents)} mental health FAQ documents")
+            for i, doc in enumerate(documents[:3]):
+                print(f"Sample FAQ {i+1}: {doc.page_content[:200]}...")
+                
+        except Exception as e:
+            print(f"Error loading mental health data: {e}")
+            raise Exception("No valid data source available. Please ensure mental health FAQ data is present or provide Kaggle credentials.")
+
+    print("Creating TF-IDF embeddings...")
+    embeddings = ImprovedTFIDFEmbeddings()
+
+    print("Creating ChromaDB vector store...")
+    client = chromadb.PersistentClient(path="./src/chroma_db")
+    
+    # Clear existing collections to prevent mixing old and new data
+    try:
+        collections = client.list_collections()
+        for collection in collections:
+            print(f"Deleting existing collection: {collection.name}")
+            client.delete_collection(collection.name)
+    except Exception as e:
+        print(f"Warning: Could not clear existing collections: {e}")
+    
+    print(f"Processing {len(documents)} documents...")
+    
+    # Check if this is a FAQ dataset and use smart retriever
+    if any("QUESTION:" in doc.page_content for doc in documents):
+        print("Using SmartFAQRetriever for better semantic matching...")
+        return SmartFAQRetriever(documents, k=4)
+    else:
+        # Use vector store for non-FAQ datasets
+        vectorstore = Chroma.from_documents(
+            documents=documents,
+            embedding=embeddings,
+            client=client
+    )
+    print("Retriever setup complete.")
+    return vectorstore.as_retriever(k=4)
+
+def setup_rag_chain() -> Runnable:
+    """Sets up the RAG chain with a prompt template and an LLM."""
+    # Define the prompt template for the LLM
+    prompt = PromptTemplate(
+        template="""You are an assistant for question-answering tasks.
+Use the following documents to answer the question.
+If you don't know the answer, just say that you don't know.
+Use three sentences maximum and keep the answer concise:
+Question: {question}
+Documents: {documents}
+Answer:
+""",
+        input_variables=["question", "documents"],
+    )
+
+    # Initialize the LLM with dolphin-llama3:8b model
+    # Note: This requires the Ollama server to be running with the specified model
+    llm = ChatOllama(
+        model="dolphin-llama3:8b",
+        temperature=0,
+    )
+
+    # Create a chain combining the prompt template and LLM
+    return prompt | llm | StrOutputParser()
+
+
+# Define the RAG application class
+class RAGApplication:
+    def __init__(self, retriever: BaseRetriever, rag_chain: Runnable):
+        self.retriever = retriever
+        self.rag_chain = rag_chain
+
+    def run(self, question: str) -> str:
+        """Runs the RAG pipeline for a given question."""
+        # Retrieve relevant documents
+        documents = self.retriever.invoke(question)
+        
+        # Debug: Print retrieved documents
+        print(f"\nDEBUG: Retrieved {len(documents)} documents for question: '{question}'")
+        for i, doc in enumerate(documents):
+            print(f"DEBUG: Document {i+1}: {doc.page_content[:200]}...")
+        
+        # Extract content from retrieved documents
+        doc_texts = "\n\n".join([doc.page_content for doc in documents])
+        
+        # Debug: Print the combined document text
+        print(f"DEBUG: Combined document text: {doc_texts[:300]}...")
+        
+        # Get the answer from the language model
+        answer = self.rag_chain.invoke({"question": question, "documents": doc_texts})
+        return answer
+
+# Main execution block
+if __name__ == "__main__":
+    load_dotenv()
+
+    # 1. Setup the components
+    retriever = setup_retriever()
+    rag_chain = setup_rag_chain()
+
+    # 2. Initialize the RAG application
+    rag_application = RAGApplication(retriever, rag_chain)
+
+    # 3. Run an example query
+    question = "What is prompt engineering"
+    print("\n--- Running RAG Application ---")
+    print(f"Question: {question}")
+    answer = rag_application.run(question)
+    print(f"Answer: {answer}")