kamkol/AB_AI_RAG_Agent

1import os2from collections import defaultdict3import tiktoken4import pickle5import shutil6import json7import time8import numpy as np9from pathlib import Path10 11from langchain_community.document_loaders import DirectoryLoader12from langchain_community.document_loaders import PyPDFLoader13from langchain_core.documents import Document14from langchain.text_splitter import RecursiveCharacterTextSplitter15from transformers import AutoModel, AutoTokenizer16import torch17import torch.nn.functional as F18from langchain_community.vectorstores import Qdrant19from qdrant_client import QdrantClient20from qdrant_client.models import Distance, VectorParams21 22from dotenv import load_dotenv23 24# Load environment variables25load_dotenv()26 27def tiktoken_len(text):28    """Count tokens using the gpt-4o-mini tokenizer"""29    tokens = tiktoken.encoding_for_model("gpt-4o-mini").encode(text)30    return len(tokens)31 32def add_page_info_to_splits(splits):33    """Process splits to add page info based on character position"""34    for split in splits:35        # Get the start position of this chunk36        start_pos = split.metadata.get("start_index", 0)37        end_pos = start_pos + len(split.page_content)38        39        # Find which page this chunk belongs to40        if "page_ranges" in split.metadata:41            for page_range in split.metadata["page_ranges"]:42                # If chunk significantly overlaps with this page range43                if (start_pos <= page_range["end"] and 44                    end_pos >= page_range["start"]):45                    # Use this page number46                    split.metadata["page"] = page_range["page"]47                    break48    return splits49 50def clean_directory(directory_path):51    """Clean a directory by removing all files and subdirectories"""52    path = Path(directory_path)53    if path.exists():54        print(f"Cleaning directory: {directory_path}")55        shutil.rmtree(path)56    57    # Wait a moment to ensure OS releases the directory handles58    time.sleep(1)59    60    path.mkdir(parents=True, exist_ok=True)61    print(f"Created clean directory: {directory_path}")62 63class ArcticEmbedder:64    def __init__(self, model_name):65        self.tokenizer = AutoTokenizer.from_pretrained(model_name)66        self.model = AutoModel.from_pretrained(model_name)67        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")68        self.model.to(self.device)69        70    def _mean_pooling(self, model_output, attention_mask):71        token_embeddings = model_output.last_hidden_state72        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()73        return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)74 75    def encode(self, texts, batch_size=32):76        all_embeddings = []77        for i in range(0, len(texts), batch_size):78            batch = texts[i:i+batch_size]79            80            encoded_input = self.tokenizer(81                batch, 82                padding=True, 83                truncation=True, 84                return_tensors="pt"85            ).to(self.device)86            87            with torch.no_grad():88                model_output = self.model(**encoded_input)89            90            batch_embeddings = self._mean_pooling(model_output, encoded_input['attention_mask'])91            batch_embeddings = F.normalize(batch_embeddings, p=2, dim=1)92            93            all_embeddings.append(batch_embeddings.cpu().numpy())94        95        return np.concatenate(all_embeddings)96 97def process_pdfs():98    """Process PDFs and create vectorstore"""99    print("Processing PDFs...")100    101    # Create processed data directory if it doesn't exist (clean it if it does)102    processed_data_dir = Path("data/processed_data")103    clean_directory(processed_data_dir)104    105    # Load all PDF documents (each page as a separate document)106    pdf_path = "notebook_version_clean/data/"107    print(f"Loading PDFs from: {pdf_path}")108    109    loader = DirectoryLoader(pdf_path, glob="*.pdf", loader_cls=PyPDFLoader)110    all_docs = loader.load()111    112    print(f"Loaded {len(all_docs)} document pages.")113    114    # Create a mapping of merged document chunks back to original pages115    docs_by_source = defaultdict(list)116    117    # Group documents by their source file118    for doc in all_docs:119        source = doc.metadata.get("source", "")120        docs_by_source[source].append(doc)121    122    # Merge pages from the same PDF but track page ranges123    merged_docs = []124    for source, source_docs in docs_by_source.items():125        # Sort by page number if available126        source_docs.sort(key=lambda x: x.metadata.get("page", 0))127        128        # Get just the filename (no path)129        filename = os.path.basename(source)130        131        # Merge the content132        merged_content = ""133        page_ranges = []134        135        for doc in source_docs:136            # Get the page number (1-indexed for human readability)137            page_num = doc.metadata.get("page", 0) + 1138            139            # Add a separator between pages for clarity140            if merged_content:141                merged_content += "\n\n"142            143            # Record where this page's content starts in the merged document144            start_pos = len(merged_content)145            merged_content += doc.page_content146            end_pos = len(merged_content)147            148            # Store the mapping of character ranges to original page numbers149            page_ranges.append({150                "start": start_pos,151                "end": end_pos,152                "page": page_num,153                "source": filename154            })155        156        # Create merged metadata that includes page mapping information157        merged_metadata = {158            "source": filename,159            "title": filename,160            "page_count": len(source_docs),161            "merged": True,162            "page_ranges": page_ranges  # Store the page ranges for later reference163        }164        165        # Create a new document with the merged content166        merged_doc = Document(page_content=merged_content, metadata=merged_metadata)167        merged_docs.append(merged_doc)168    169    print(f"Created {len(merged_docs)} merged documents.")170    171    # Split documents172    text_splitter = RecursiveCharacterTextSplitter(173        chunk_size=150,174        chunk_overlap=100,175        length_function=tiktoken_len,176        add_start_index=True177    )178    179    # Split and then process to add page information180    raw_splits = text_splitter.split_documents(merged_docs)181    split_chunks = add_page_info_to_splits(raw_splits)182    183    print(f"Created {len(split_chunks)} chunks.")184    185    # Save chunks for later use186    with open(processed_data_dir / "chunks.pkl", "wb") as f:187        pickle.dump(split_chunks, f)188    189    190    191    # Initialize custom embedding model192    try:193        embedding_model = ArcticEmbedder("kamkol/ab_testing_finetuned_arctic_ft-36dfff22-0696-40d2-b3bf-268fe2ff2aec")194        print("Successfully loaded ArcticEmbedder model")195    196    except Exception as e:197        print(f"Error loading model: {str(e)}")198        raise RuntimeError(f"Error initializing SentenceTransformer model: {str(e)}")199    200    print("Embedding document chunks (this may take a while)...")201    # Create a dictionary to store documents and their embeddings202    embedded_docs = []203    204    # Embed in batches to avoid API rate limits205    batch_size = 50206    for i in range(0, len(split_chunks), batch_size):207        batch = split_chunks[i:i+batch_size]208        209        # Extract text210        texts = [doc.page_content for doc in batch]211        212        # Get embeddings213        embeddings = embedding_model.encode(texts)214        215        # Store with metadata216        for j, doc in enumerate(batch):217            embedded_docs.append({218                "id": i + j,219                "text": doc.page_content,220                "metadata": doc.metadata,221                "embedding": embeddings[j]222            })223        224        # Print progress225        print(f"Embedded {min(i+batch_size, len(split_chunks))}/{len(split_chunks)} chunks")226    227    # Save the embedded docs for later use228    with open(processed_data_dir / "embedded_docs.pkl", "wb") as f:229        pickle.dump(embedded_docs, f)230    231    print("Processing complete. All data saved to data/processed_data/")232 233if __name__ == "__main__":234    process_pdfs()