kchen707/wedding-bundle-builder

1#!/usr/bin/env python32"""3One-time embedding cache builder.4 5Run this LOCALLY once with your OpenRouter API key set. It will:6  1. Load and preprocess the vendor xlsx (same logic as the app)7  2. Compute embeddings for every vendor in batches8  3. Save them to data/vendor_embeddings_<fingerprint>.pkl9 10Then commit the resulting .pkl file to the Space repo. The deployed11app loads it at startup so it never has to call the embedding API12for vendor data — only for user queries.13 14Usage15-----16   $ export OPENROUTER_API_KEY=sk-or-v1-...17   $ python scripts/build_embeddings.py18 19Cost estimate20-------------211,100 vendors × ~400 tokens each × $0.13/1M tokens (text-embedding-3-large)22≈ $0.06 per full rebuild. Cheap enough to re-run if your dataset changes.23"""24import os25import sys26import pickle27import time28from pathlib import Path29 30# Make the parent directory importable so `from config import ...` works31# regardless of where this script is invoked from.32HERE = Path(__file__).resolve().parent33ROOT = HERE.parent34sys.path.insert(0, str(ROOT))35 36import numpy as np  # noqa: E40237from openai import OpenAI  # noqa: E40238 39from config import EMBEDDING_MODEL  # noqa: E40240from data_loader import (  # noqa: E40241    load_vendor_dataframe,42    corpus_fingerprint,43    DATA_DIR,44)45 46 47def main():48    api_key = os.environ.get("OPENROUTER_API_KEY", "").strip()49    if not api_key:50        print("ERROR: Set OPENROUTER_API_KEY environment variable first.")51        print("  e.g.  export OPENROUTER_API_KEY=sk-or-v1-...")52        sys.exit(1)53 54    print("Loading and preprocessing vendor data...")55    df = load_vendor_dataframe()56    texts = df["embedding_text"].tolist()57    print(f"  {len(df)} active vendors, "58          f"avg embedding_text length = {int(np.mean([len(t) for t in texts]))} chars")59 60    fingerprint = corpus_fingerprint(texts, EMBEDDING_MODEL)61    out_path = DATA_DIR / f"vendor_embeddings_{fingerprint}.pkl"62    print(f"  fingerprint: {fingerprint}")63    print(f"  output:      {out_path}")64 65    if out_path.exists():66        print("\n✓ Cache already exists with this fingerprint. Nothing to do.")67        return68 69    print(f"\nGenerating embeddings with {EMBEDDING_MODEL}...")70    client = OpenAI(base_url="https://openrouter.ai/api/v1", api_key=api_key)71 72    batch_size = 10073    total_batches = (len(texts) + batch_size - 1) // batch_size74    all_embeddings = []75 76    for i in range(0, len(texts), batch_size):77        batch = texts[i:i + batch_size]78        batch_num = i // batch_size + 179        print(f"  Batch {batch_num}/{total_batches} ({len(batch)} vendors)...",80              end=" ", flush=True)81 82        resp = client.embeddings.create(input=batch, model=EMBEDDING_MODEL)83        all_embeddings.extend(item.embedding for item in resp.data)84        print("✓")85 86        if batch_num < total_batches:87            time.sleep(0.5)  # gentle on rate limits88 89    embeddings = np.array(all_embeddings, dtype=np.float32)90 91    DATA_DIR.mkdir(parents=True, exist_ok=True)92    with open(out_path, "wb") as f:93        pickle.dump({94            "embeddings":  embeddings,95            "count":       len(df),96            "fingerprint": fingerprint,97            "model":       EMBEDDING_MODEL,98        }, f)99 100    size_mb = out_path.stat().st_size / (1024 * 1024)101    print(f"\n✅ Saved {embeddings.shape} → {out_path.name} ({size_mb:.1f} MB)")102    print(f"   Commit this file to the repo so the deployed app can use it.")103 104 105if __name__ == "__main__":106    main()107