import logging import math import sys import os import re import numpy as np import pandas as pd from tqdm import tqdm import ast from collections import defaultdict from datasets import load_dataset, Dataset from langchain_community.vectorstores import FAISS from langchain_community.embeddings import HuggingFaceEmbeddings from langchain_text_splitters import RecursiveCharacterTextSplitter # BM25 for lexical search from rank_bm25 import BM25Okapi import nltk nltk.download('punkt', quiet=True) nltk.download('punkt_tab', quiet=True) from nltk.tokenize import word_tokenize from langchain_core.documents import Document logging.basicConfig( level=logging.INFO, format='%(asctime)s | %(levelname)s | %(message)s', datefmt='%H:%M:%S', handlers=[ logging.StreamHandler(sys.stdout) # Try console too ] ) logger = logging.getLogger(__name__) # Embedding model EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2" # Hybrid retrieval weights SEMANTIC_WEIGHT = 0.4 BM25_WEIGHT = 0.3 ENTITY_WEIGHT = 0.3 # Retrieval parameters CHUNK_SIZE = 256 CHUNK_OVERLAP = 100 DEFAULT_K = 10 DEFAULT_FETCH_K = 10000 # Dataset settings DATASET_USERS_NAME = "srirxml/PANORAMA-Plus" DATASET_TEXTS_NAME = "srirxml/PANORAMA" SPLIT = "train" MAX_USERS = 1000 MAX_TEXTS_PER_USER = None # Column mappings USER_ID_COL_USERS = "Unique ID" USER_ID_COL_TEXTS = "id" TEXT_COL = "text" MIN_CHARS = 10 # Locale-to-location mapping LOCALE_TO_LOCATION = { "en_PH": "Philippines", "en_CA": "Canada", "en_US": "United States", "en_IE": "Ireland", "en_NZ": "New Zealand", "en_IN": "India", "en_AU": "Australia", "en_GB": "United Kingdom", "en_IL": "Israel", "en_DE": "Germany", "en_IT": "Italy", "en_FR": "France", } # Sensitive attributes for privacy analysis SENSITIVE_ATTRIBUTES = ["Age bin", "Gender", "Marital Status", "Finance Status", "Education", "Locale"] ATTRIBUTE_VALUES_MAP = { "Gender": ["Female", "Male"], "Age bin": ["0-17", "18-29", "30-44", "45-59", "60+"], "Marital Status": ["Single", "Married", "Divorced", "Widowed"], "Finance Status": ["Low", "Medium", "High"], "Locale": [LOCALE_TO_LOCATION["en_PH"], LOCALE_TO_LOCATION["en_CA"], LOCALE_TO_LOCATION["en_US"], LOCALE_TO_LOCATION["en_IE"], LOCALE_TO_LOCATION["en_NZ"], LOCALE_TO_LOCATION["en_IN"], LOCALE_TO_LOCATION["en_AU"], LOCALE_TO_LOCATION["en_GB"], LOCALE_TO_LOCATION["en_IL"], LOCALE_TO_LOCATION["en_DE"], LOCALE_TO_LOCATION["en_IT"], LOCALE_TO_LOCATION["en_FR"]], "Education": ["High School", "Bachelor's", "Master's", "PhD"] } def strip_special_chars(text): """Strip digits and special characters, keeping only letters.""" return re.sub(r'[^a-zA-Z]', '', text.lower()) def age_to_bin(age): if pd.isna(age): return pd.NA try: age = int(age) if age < 18: return "0-17" elif age < 30: return "18-29" elif age < 45: return "30-44" elif age < 60: return "45-59" else: return "60+" except: return pd.NA def safe_parse_handles(x): """Parse social media handles from stringified dict.""" if x is None or (isinstance(x, float) and pd.isna(x)): return {} s = str(x).strip() if not s or s.lower() == "nan": return {} try: v = ast.literal_eval(s) return v if isinstance(v, dict) else {} except Exception: return {} def build_persona(row): """Build a persona string from user profile.""" first = str(row.get("First Name", "") or "").strip() last = str(row.get("Last Name", "") or "").strip() handles_dict = safe_parse_handles(row.get("Social Media Handles")) handles = [str(v).strip() for v in handles_dict.values() if v and str(v).strip()] name_part = (first + " " + last).strip() handle_part = ", ".join(handles) if name_part and handle_part: return f"{name_part}; {handle_part}" elif name_part: return name_part elif handle_part: return handle_part else: return str(row.get(USER_ID_COL_USERS, "")).strip() def load_panorama_data(): """Load and merge PANORAMA datasets.""" print("\n" + "=" * 80) print("LOADING DATA") print("=" * 80) ds_users = load_dataset(DATASET_USERS_NAME, split=SPLIT) ds_texts = load_dataset(DATASET_TEXTS_NAME, split=SPLIT) users_df = ds_users.to_pandas() texts_df = ds_texts.to_pandas() # Select first N users first_user_ids = ( users_df[USER_ID_COL_USERS] .dropna() .astype(str) .drop_duplicates() .head(MAX_USERS) .tolist() ) users_df_small = users_df[users_df[USER_ID_COL_USERS].astype(str).isin(first_user_ids)].copy() texts_df_small = texts_df[texts_df[USER_ID_COL_TEXTS].astype(str).isin(first_user_ids)].copy() # Clean texts texts_df_small[TEXT_COL] = texts_df_small[TEXT_COL].astype(str).str.strip() texts_df_small = texts_df_small[texts_df_small[TEXT_COL].str.len() >= MIN_CHARS].copy() texts_df_small = texts_df_small.drop_duplicates(subset=[USER_ID_COL_TEXTS, TEXT_COL]).copy() # Optional: cap texts per user if MAX_TEXTS_PER_USER is not None: texts_df_small = ( texts_df_small .groupby(USER_ID_COL_TEXTS, as_index=False, sort=False) .head(int(MAX_TEXTS_PER_USER)) .copy() ) users_df_small["Age bin"] = users_df_small["Age"].apply(age_to_bin) users_df_small["Finance Status"] = users_df_small["Finance Status"].apply( lambda x: "High" if "high" in x.lower() else ( "Medium" if "medium" in x.lower() else ("Low" if "low" in x.lower() else pd.NA)) ) users_df_small["Education"] = users_df_small["Education Info"].apply( lambda x: "Bachelor's" if x in ["Bachelor's", "Some College", "Diploma", "Associate's", "Professional Certificate", "Vocational Training"] else ("High School" if x in ["Less than High School", "High School"] else x) ) users_df_small["Locale"] = users_df_small["Locale"].apply( lambda x: LOCALE_TO_LOCATION[x] if x in LOCALE_TO_LOCATION else x) # Merge merged = users_df_small.merge( texts_df_small, left_on=USER_ID_COL_USERS, right_on=USER_ID_COL_TEXTS, how="left", ) print(f"✓ Users selected: {len(first_user_ids)}") print(f"✓ Texts after per-user dedup: {len(texts_df_small)}") print(f"✓ Merged rows (users x texts): {len(merged)}") return merged, users_df_small class HybridRetriever: """ Hybrid retriever combining semantic (FAISS), lexical (BM25), and entity-based search with Reciprocal Rank Fusion (RRF). The ``retrieve`` method accepts an optional ``min_similarity`` threshold. When set, only documents whose normalised semantic similarity to the query meets or exceeds that value are eligible for retrieval. This makes the system sensitive to Input-DP perturbation: a heavily noised query will drift away from the persona corpus and return fewer — or zero — documents, so the system's output correctly reflects the privacy protection that is active. """ def __init__( self, documents, embedding_model=EMBEDDING_MODEL, semantic_weight=SEMANTIC_WEIGHT, bm25_weight=BM25_WEIGHT, entity_weight=ENTITY_WEIGHT, chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP, ): self.semantic_weight = semantic_weight self.bm25_weight = bm25_weight self.entity_weight = entity_weight # Text splitter self.splitter = RecursiveCharacterTextSplitter( chunk_size=chunk_size, chunk_overlap=chunk_overlap, length_function=len, ) print("Building retriever components...") # Build FAISS index for semantic search self.embeddings = HuggingFaceEmbeddings(model_name=embedding_model) self.vectorstore = FAISS.from_documents(documents, self.embeddings) print(f" ✓ FAISS index built ({len(documents)} docs)") # Build BM25 index for lexical search self.bm25_corpus = [doc.page_content for doc in documents] tokenized_corpus = [word_tokenize(doc.lower()) for doc in self.bm25_corpus] self.bm25 = BM25Okapi(tokenized_corpus) print(f" ✓ BM25 index built") # Store documents for entity matching self.documents = documents # Extract entities (usernames, identifiers) from metadata # Strip digits and special characters for better matching self.entity_index = defaultdict(list) for i, doc in tqdm(enumerate(documents)): persona = doc.metadata.get("persona", "") if persona: # Extract potential identifiers and strip special chars tokens = re.split(r'[;,\s]+', persona.lower()) for token in tokens: # Strip special characters and digits clean_token = strip_special_chars(token) if len(clean_token) > 2: # Skip very short tokens self.entity_index[clean_token].append(i) print(f" ✓ Entity index built ({len(self.entity_index)} unique entities)") # ── Content-to-corpus-index lookup (used by threshold filtering) ────── # Maps page_content → list of corpus indices so that FAISS results # (which are Document objects, not indices) can be mapped back to # their position in self.documents. Duplicate page_content values # are handled by storing all matching indices. self._content_to_indices = defaultdict(list) for i, doc in enumerate(self.documents): self._content_to_indices[doc.page_content].append(i) def _reciprocal_rank_fusion(self, rankings, k=60): """Combine multiple rankings using RRF.""" scores = defaultdict(float) sources = defaultdict(dict) for source_name, ranking in rankings.items(): for rank, doc_id in enumerate(ranking, start=1): scores[doc_id] += 1.0 / (k + rank) sources[doc_id][source_name] = rank return scores, sources def retrieve_semantic_only(self, query, k=10): """Retrieve using only semantic search (for comparison).""" return self.vectorstore.similarity_search(query, k=k) def retrieve(self, query, k=10, fetch_k=100, min_similarity=None): """Hybrid retrieval combining semantic, BM25, and entity matching. Parameters ---------- query : str The search query (may be DP-perturbed). k : int Maximum number of documents to return. fetch_k : int Candidate pool size passed to FAISS. min_similarity : float or None When set, only documents whose normalised semantic similarity to the query meets or exceeds this value are eligible. Documents below the threshold are excluded from ALL three ranking components (semantic, BM25, entity) before RRF. Pass None (default) to disable filtering and preserve the original behaviour. Returns ------- list[Document] Up to k documents, potentially fewer (or empty) when min_similarity is strict relative to the query. """ # ── 1. Semantic search WITH scores ─────────────────────────────────── try: candidates = self.vectorstore.similarity_search_with_score( query, k=fetch_k ) except Exception: # Fallback: vectorstore does not expose scores → no threshold docs = self.vectorstore.similarity_search(query, k=fetch_k) candidates = [(doc, 0.0) for doc in docs] if not candidates: return [] raw_docs, raw_scores = zip(*candidates) raw_scores = list(raw_scores) # ── Normalise scores to similarity ∈ [0, 1] ───────────────────────── # FAISS/IndexFlatL2 returns L2 distances (ascending: closer = smaller). # IndexFlatIP on normalised vectors returns cosine scores (descending). if len(raw_scores) >= 2 and raw_scores[0] <= raw_scores[-1]: # L2 distances: invert so that higher = more similar max_dist = max(raw_scores) + 1e-10 similarities = [1.0 - s / max_dist for s in raw_scores] else: # Already similarity scores; clamp to [0, 1] similarities = [max(0.0, min(1.0, s)) for s in raw_scores] # ── Apply similarity threshold ──────────────────────────────────────── if min_similarity is not None: passing_pairs = [ (doc, sim) for doc, sim in zip(raw_docs, similarities) if sim >= min_similarity ] if not passing_pairs: logger.info( " HybridRetriever: 0/%d candidates above threshold %.3f" " — returning []", len(raw_docs), min_similarity, ) return [] # Build the set of *corpus* indices that passed the threshold. # This is used to restrict BM25 and entity rankings to the # same document subset, ensuring all three components only # vote for threshold-passing documents. valid_corpus_indices = set() for doc, _ in passing_pairs: for idx in self._content_to_indices.get(doc.page_content, []): valid_corpus_indices.add(idx) # Semantic ranking: use actual corpus indices for filtered docs # so they are consistent with BM25/entity namespace. semantic_ranking = [] for doc, _ in passing_pairs: for idx in self._content_to_indices.get(doc.page_content, []): semantic_ranking.append(idx) break # one representative index per doc is enough for RRF else: # No threshold: preserve original behaviour (range-based indices). valid_corpus_indices = None semantic_ranking = list(range(len(raw_docs))) # ── 2. BM25 search ─────────────────────────────────────────────────── query_tokens = word_tokenize(query.lower()) bm25_scores = self.bm25.get_scores(query_tokens) bm25_ranking = np.argsort(bm25_scores)[::-1][:fetch_k].tolist() if valid_corpus_indices is not None: bm25_ranking = [i for i in bm25_ranking if i in valid_corpus_indices] # ── 3. Entity matching ─────────────────────────────────────────────── query_lower = query.lower() entity_matches = set() for query_token in re.split(r'[;,\s.!?]+', query_lower): clean_query_token = strip_special_chars(query_token) if len(clean_query_token) > 2: for entity, doc_ids in self.entity_index.items(): if clean_query_token in entity or entity in clean_query_token: entity_matches.update(doc_ids) entity_ranking = list(entity_matches)[:fetch_k] if valid_corpus_indices is not None: entity_ranking = [i for i in entity_ranking if i in valid_corpus_indices] # ── 4. Reciprocal Rank Fusion ───────────────────────────────────────── weighted_rankings = {} if self.semantic_weight > 0: weighted_rankings["semantic"] = semantic_ranking if self.bm25_weight > 0: weighted_rankings["bm25"] = bm25_ranking if self.entity_weight > 0 and entity_ranking: weighted_rankings["entity"] = entity_ranking scores, sources = self._reciprocal_rank_fusion(weighted_rankings) sorted_doc_ids = sorted( scores.keys(), key=lambda x: scores[x], reverse=True )[:k] result = [self.documents[i] for i in sorted_doc_ids] logger.info( " HybridRetriever: returning %d/%d docs (min_similarity=%s)", len(result), len(candidates), f"{min_similarity:.3f}" if min_similarity is not None else "None", ) return result # ============================================================================== # DIFFERENTIAL PRIVACY RETRIEVER # ============================================================================== class DPRetriever: """Differentially private retriever using the Exponential Mechanism. Sensitivity is computed *empirically* from the actual utility range of the candidate pool rather than a fixed constant. This data-adaptive approach (Dwork & Roth 2014; Koga et al. 2024) gives a tighter bound, improving the privacy-utility trade-off without weakening the formal ε-DP guarantee. For the exponential mechanism the sensitivity Δu of utility function u is: Δu = max_{d, d'} |u(d, q) − u(d', q)| When u(d, q) = max_dist − dist(d, q) (distance-to-similarity inversion), Δu equals the range of utilities across the candidate pool. We clamp it to ``sensitivity_cap`` to guard against degenerate cases where all candidates are equidistant from the query. Parameters ---------- base_retriever : HybridRetriever The underlying retriever whose vectorstore provides FAISS scores. epsilon : float DP privacy budget. Smaller ε → stronger privacy. sensitivity_cap : float Floor value for the empirical sensitivity (default 1e-3). For most corpora the empirical range will be far larger than this cap so it has no practical effect. """ def __init__(self, base_retriever, epsilon=1.0, sensitivity_cap=1e-3): self.base_retriever = base_retriever self.epsilon = epsilon self.sensitivity_cap = sensitivity_cap def retrieve(self, query, k=10, fetch_k=100, seed=None): """Retrieve up to k documents with differential privacy. Uses the Exponential Mechanism: each candidate document is sampled with probability proportional to exp(ε · u(d) / 2Δu), where u is the normalised similarity utility and Δu is the empirical range. Parameters ---------- query : str Search query (may be DP-perturbed). k : int Number of documents to return. fetch_k : int Candidate pool size fetched from FAISS before sampling. seed : int or None Optional random seed for reproducibility. Returns ------- list[Document] k documents sampled with DP-weighted probabilities. """ if seed is not None: np.random.seed(seed) # Fetch candidate pool try: candidates = self.base_retriever.vectorstore.similarity_search_with_score( query, k=fetch_k ) except Exception: # Fallback: vectorstore does not expose scores docs = self.base_retriever.vectorstore.similarity_search(query, k=fetch_k) return docs[:k] if not candidates: return [] docs, base_scores = zip(*candidates) base_scores = np.array(base_scores) # ── Convert raw scores to a utility where higher = better ──────────── # FAISS with L2: scores are distances (ascending) → invert # FAISS with IP: scores are similarities (descending) → use as-is if base_scores[0] >= base_scores[-1]: utilities = base_scores.copy() # already similarity scores else: max_dist = base_scores.max() + 1e-10 utilities = max_dist - base_scores # invert L2 distances # ── Empirical sensitivity (Dwork & Roth 2014; Koga et al. 2024) ───── u_range = float(utilities.max() - utilities.min()) sensitivity = max(u_range, self.sensitivity_cap) # ── Exponential mechanism: P(d) ∝ exp(ε · u(d) / 2Δu) ────────────── probabilities = np.exp((self.epsilon * utilities) / (2.0 * sensitivity)) probabilities = probabilities / probabilities.sum() # Sample k documents without replacement selected_indices = np.random.choice( len(docs), size=min(k, len(docs)), replace=False, p=probabilities ) return [docs[i] for i in selected_indices] def save_retriever_components(retriever, save_path): """ Save retriever components (documents + config) instead of entire object. Avoids FAISS serialization issues. """ import pickle import os logger.info(f"💾 Saving retriever components to {save_path}...") # Extract all components needed to rebuild the retriever components = { 'documents': retriever.documents, 'config': { 'embedding_model': retriever.embeddings.model_name, 'semantic_weight': retriever.semantic_weight, 'bm25_weight': retriever.bm25_weight, 'entity_weight': retriever.entity_weight, 'chunk_size': retriever.splitter._chunk_size, 'chunk_overlap': retriever.splitter._chunk_overlap, }, 'metadata': { 'num_documents': len(retriever.documents), 'num_entities': len(retriever.entity_index), } } # Save to pickle os.makedirs(os.path.dirname(save_path) if os.path.dirname(save_path) else '.', exist_ok=True) with open(save_path, 'wb') as f: pickle.dump(components, f, protocol=pickle.HIGHEST_PROTOCOL) logger.info(f" ✓ Saved {len(components['documents'])} documents") logger.info(f"✅ Retriever components saved to {save_path}") def load_retriever_components(load_path): """ Load components and rebuild HybridRetriever from scratch. Rebuilds FAISS, BM25, and entity indices in current environment. """ import pickle logger.info(f"📁 Loading retriever components from {load_path}...") # Load components with open(load_path, 'rb') as f: components = pickle.load(f) documents = components['documents'] config = components['config'] logger.info(f" ✓ Loaded {len(documents)} documents") logger.info(f" ✓ Config: semantic_weight={config['semantic_weight']}, " f"bm25_weight={config['bm25_weight']}, entity_weight={config['entity_weight']}") # Rebuild the retriever from scratch using current environment's packages logger.info("🔨 Rebuilding retriever indices (this takes ~30-60 seconds)...") retriever = HybridRetriever( documents=documents, embedding_model=config['embedding_model'], semantic_weight=config['semantic_weight'], bm25_weight=config['bm25_weight'], entity_weight=config['entity_weight'], chunk_size=config['chunk_size'], chunk_overlap=config['chunk_overlap'], ) logger.info("✅ Retriever rebuilt successfully") return retriever def build_documents(merged, users_df): """Build documents and metadata for RAG.""" print("\n" + "=" * 80) print("BUILDING DOCUMENTS") print("=" * 80) texts = [] metas = [] user_meta_cols = [c for c in users_df.columns if c != USER_ID_COL_USERS] for _, row in tqdm(merged.iterrows(), total=len(merged), desc="Building documents"): t = row.get(TEXT_COL, "") if not isinstance(t, str): continue t = t.strip() if len(t) < MIN_CHARS: continue persona = build_persona(row) texts.append(t) meta = {c: row.get(c) for c in user_meta_cols} meta["user_id"] = str(row.get(USER_ID_COL_USERS)) meta["persona"] = persona if "source" in merged.columns: meta["source"] = row.get("source") if "content_type" in merged.columns: meta["content_type"] = row.get("content_type") metas.append(meta) print(f"✓ Documents created: {len(texts)}") return texts, metas if __name__ == "__main__": # retriever = load_retriever_components("./faiss_persona_sarah_chen_retriever_components.pkl") retriever = load_retriever_components("./faiss_panorama_retriever_components.pkl") # save_retriever_components(retriever, "./faiss_persona_sarah_chen_retriever_components.pkl") save_retriever_components(retriever, "./faiss_panorama_retriever_components.pkl") RETRIEVER_SAVE_PATH = r"C:\Users\user\Datasets\Panorama Synthetic Data RAG\faiss_panorama_retriever_components.pkl" if os.path.exists(RETRIEVER_SAVE_PATH): retriever = load_retriever_components(RETRIEVER_SAVE_PATH) print(retriever.retrieve("Hi, I am Raymond Phillips", 10)[0:5]) else: merged, users_df = load_panorama_data() # ------------------------------------------------------------------------- # STEP 2: Build Documents # ------------------------------------------------------------------------- texts, metas = build_documents(merged, users_df) # Create Document objects documents = [ Document(page_content=text, metadata=meta) for text, meta in zip(texts, metas) ] # Saving (in your local environment): retriever = HybridRetriever( documents, embedding_model=EMBEDDING_MODEL, semantic_weight=SEMANTIC_WEIGHT, bm25_weight=BM25_WEIGHT, entity_weight=ENTITY_WEIGHT, chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP, ) save_retriever_components(retriever, RETRIEVER_SAVE_PATH)