create_rag_retriever.py

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)