kchen707/wedding-bundle-builder

1"""2Data loading module: reads the local vendor xlsx, runs the same3preprocessing as the notebook (active-only filter, type coercion,4enriched embedding_text), and loads the pre-computed embedding cache.5 6If the embedding cache is missing, this module will raise a clear7error pointing the operator at scripts/build_embeddings.py.8"""9 10from pathlib import Path11import hashlib12import pickle13 14import numpy as np15import pandas as pd16 17from config import EMBEDDING_MODEL18 19 20DATA_DIR = Path(__file__).parent / "data"21VENDOR_FILE = DATA_DIR / "SD_Wedding_Vendors_Combined.xlsx"22EMBEDDING_CACHE_GLOB = "vendor_embeddings_*.pkl"23 24 25# ----------------------------------------------------------------------26# Preprocessing helpers (lifted from cell 7)27# ----------------------------------------------------------------------28 29def _safe_str(v) -> str:30    """Convert a cell to a clean string, handling NaN. (cell 7)"""31    try:32        if v is None or pd.isna(v):33            return ""34    except Exception:35        pass36    return str(v).strip()37 38 39def build_enriched_embedding_text(row) -> str:40    """41    Construct a richer embedding text that includes review signals.42    (cell 7)43    """44    parts = []45 46    base = _safe_str(row.get("embedding_text")) or _safe_str(row.get("description"))47    if base:48        parts.append(base)49 50    style = _safe_str(row.get("style_keywords"))51    if style:52        parts.append(f"Style: {style}.")53 54    specialty = _safe_str(row.get("specialty_tags"))55    if specialty:56        parts.append(f"Specialties: {specialty}.")57 58    cultural = _safe_str(row.get("cultural_expertise"))59    if cultural and cultural.lower() not in ("nan", "none", ""):60        parts.append(f"Cultural expertise: {cultural}.")61 62    reviews = _safe_str(row.get("sample_review_snippets"))63    if reviews:64        parts.append(f"What couples say: {reviews}")65 66    services = _safe_str(row.get("services_offered"))67    if services and len(services) > 10:68        parts.append(f"Services: {services[:300]}.")69 70    return " ".join(parts)71 72 73# ----------------------------------------------------------------------74# Cache fingerprinting (matches cell 9)75# ----------------------------------------------------------------------76 77def corpus_fingerprint(texts, model: str) -> str:78    """16-char hex fingerprint of (model, ordered corpus). Identical to cell 9."""79    h = hashlib.sha256()80    h.update(model.encode("utf-8"))81    for t in texts:82        h.update(b"\x00")83        h.update(t.encode("utf-8", errors="ignore"))84    return h.hexdigest()[:16]85 86 87# ----------------------------------------------------------------------88# Public loaders89# ----------------------------------------------------------------------90 91def load_vendor_dataframe() -> pd.DataFrame:92    """93    Load the vendor xlsx and apply cell-7 preprocessing.94    Returns a dataframe whose row order is stable across runs as long as95    the source xlsx is unchanged.96    """97    if not VENDOR_FILE.exists():98        raise FileNotFoundError(99            f"Vendor data not found at {VENDOR_FILE}. "100            f"Make sure SD_Wedding_Vendors_Combined.xlsx is in the data/ folder."101        )102 103    df = pd.read_excel(VENDOR_FILE)104 105    # cell 7: active vendors only106    df = df[df["is_active"] == True].copy()107 108    # cell 7: numeric coercion109    for col in ["price_min", "price_max", "guest_capacity_min", "guest_capacity_max",110                "avg_rating", "num_reviews"]:111        df[col] = pd.to_numeric(df[col], errors="coerce")112 113    # cell 7: enriched embedding text114    df["embedding_text"] = df.apply(build_enriched_embedding_text, axis=1)115    df = df[df["embedding_text"].str.len() > 20].copy()116 117    # Reset the index so positional lookups (used in semantic_search_bundle)118    # match the embedding row order. The notebook relied on index continuity119    # post-filter; we make that explicit here.120    df = df.reset_index(drop=True)121 122    return df123 124 125def load_embeddings(df: pd.DataFrame) -> np.ndarray:126    """127    Load pre-computed embeddings from disk. The cache file is keyed128    by (model, corpus fingerprint) so we can detect drift.129 130    Returns a (len(df), EMBEDDING_DIM) numpy array.131    """132    texts = df["embedding_text"].tolist()133    expected_fp = corpus_fingerprint(texts, EMBEDDING_MODEL)134    expected_path = DATA_DIR / f"vendor_embeddings_{expected_fp}.pkl"135 136    # Try the exact-fingerprint file first (fast path)137    if expected_path.exists():138        with open(expected_path, "rb") as f:139            cache = pickle.load(f)140        embeddings = cache["embeddings"]141        if (cache.get("fingerprint") == expected_fp142                and len(embeddings) == len(df)):143            return embeddings144        # fall through to looser match if fingerprint somehow mismatches145 146    # Fallback: scan the data folder for any cache pickle and accept the147    # first one whose row count matches. This makes the deployment148    # forgiving if you committed a cache built on a slightly different149    # run of the preprocessing.150    for cache_path in DATA_DIR.glob(EMBEDDING_CACHE_GLOB):151        try:152            with open(cache_path, "rb") as f:153                cache = pickle.load(f)154            embeddings = cache["embeddings"]155            if len(embeddings) == len(df):156                print(f"⚠️  Using cache {cache_path.name} (row count matches "157                      f"but fingerprint differs from current preprocessing).")158                return embeddings159        except Exception:160            continue161 162    # Nothing usable found163    raise FileNotFoundError(164        f"\n\nNo embedding cache found in {DATA_DIR}.\n\n"165        f"Run the build script once locally to generate one:\n"166        f"   export OPENROUTER_API_KEY=sk-or-...\n"167        f"   python scripts/build_embeddings.py\n\n"168        f"This will create data/vendor_embeddings_<fingerprint>.pkl which you\n"169        f"then commit to the Space repo. Cost is roughly $0.50–$1 once.\n"170    )171