image-classifier/build_code_embeddings.py

"""
build_code_embeddings.py — Embed USPTO Design Codes for RAG retrieval
=======================================================================
One-time script that reads uspto_design_codes.json (from the scraper) and
embeds every code's description using Voyage AI. The output feeds the
RAG layer of design_code_classifier.py.

WHY THIS EXISTS:
The classifier needs to constrain Claude to USPTO's actual code vocabulary
(no hallucinating fake codes). Embedding every code's description once
lets us do fast vector search at classification time: Claude describes
the image, we retrieve the most semantically similar codes, Claude picks
from that validated menu.

SETUP:
    pip install voyageai numpy python-dotenv

ENV VARS:
    VOYAGE_API_KEY — Voyage AI API key (free tier covers 200M tokens)

USAGE:
    python build_code_embeddings.py

    # Re-embed with a different model:
    python build_code_embeddings.py --model voyage-3-large

OUTPUT:
    uspto_code_embeddings.pkl — pickle dict with:
      - codes:        list[str]     ordered list of XX.YY.ZZ codes
      - descriptions: list[str]     parallel list of code descriptions
      - categories:   list[str]     parent category for each code (XX)
      - embeddings:   np.ndarray    shape (N, dim), float32
      - metadata:     dict          model name, dim, timestamp

COST:
    ~1,300 codes × ~15 tokens each = ~20K tokens.
    voyage-3.5 is ~$0.06 per 1M tokens → effectively free under the
    200M-token free tier.
"""

import os
import sys
import json
import pickle
import logging
from datetime import datetime, timezone
from pathlib import Path
from typing import List

import numpy as np
from dotenv import load_dotenv

try:
    import voyageai
except ImportError:
    print("ERROR: Voyage AI SDK not installed. Run:")
    print("  pip install voyageai")
    sys.exit(1)


# ============================================================================
# CONFIG
# ============================================================================

env_path = Path(__file__).parent / ".env"
load_dotenv(dotenv_path=env_path)

VOYAGE_API_KEY = os.getenv("VOYAGE_API_KEY")

DEFAULT_MODEL = "voyage-3.5"   # current general-purpose default; voyage-3-large is the premium upgrade
DEFAULT_DIM = 1024              # default for voyage-3.5; do not change without re-embedding

# Voyage allows up to 1,000 texts per batch; we have ~1,300 codes, so 2 batches
BATCH_SIZE = 1000

INPUT_PATH = Path(__file__).parent / "uspto_design_codes.json"
OUTPUT_PATH = Path(__file__).parent / "uspto_code_embeddings.pkl"

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
)
logger = logging.getLogger("embeddings")


# ============================================================================
# LOAD CODES FROM SCRAPED MANUAL
# ============================================================================

def load_codes() -> tuple[List[str], List[str], List[str]]:
    """Read uspto_design_codes.json and flatten into parallel lists.

    Returns:
        (codes, descriptions, categories) — three lists of equal length,
        where codes[i] corresponds to descriptions[i] from category categories[i].
    """
    if not INPUT_PATH.exists():
        logger.error(f"❌ {INPUT_PATH} not found. Run scrape_uspto_design_codes.py first.")
        sys.exit(1)

    data = json.loads(INPUT_PATH.read_text())
    codes_dict = data.get("categories", {})

    codes: List[str] = []
    descriptions: List[str] = []
    categories: List[str] = []
    seen = set()  # deduplicate (the scraper produced some cross-listed codes)

    for category_id, category_data in codes_dict.items():
        for division_id, division_data in category_data.get("divisions", {}).items():
            for section_code, section_data in division_data.get("sections", {}).items():
                if section_code in seen:
                    continue
                description = section_data.get("description", "").strip()
                if not description:
                    continue
                codes.append(section_code)
                descriptions.append(description)
                # Use the section's *real* category prefix (XX), not the JSON parent
                # — this self-corrects the scraper's cross-listing duplicates
                categories.append(section_code.split(".")[0])
                seen.add(section_code)

    logger.info(f"📋 Loaded {len(codes)} unique codes from {INPUT_PATH.name}")
    return codes, descriptions, categories


# ============================================================================
# EMBED VIA VOYAGE
# ============================================================================

def build_searchable_text(code: str, description: str) -> str:
    """Construct the text that gets embedded for each code.

    We include the code itself in the text — the digits give the embedding
    a tiny extra signal of which category/division things belong to, which
    helps when descriptions are very generic ("Other plants" appears in
    multiple divisions and would otherwise be indistinguishable).
    """
    return f"USPTO Design Code {code}: {description}"


def embed_descriptions(
    descriptions_with_codes: List[str],
    model: str,
) -> np.ndarray:
    """Call Voyage to embed all descriptions. Returns (N, dim) float32 array."""
    if not VOYAGE_API_KEY:
        logger.error("❌ VOYAGE_API_KEY not set in .env")
        sys.exit(1)

    client = voyageai.Client(api_key=VOYAGE_API_KEY)

    all_embeddings: List[List[float]] = []

    for batch_idx in range(0, len(descriptions_with_codes), BATCH_SIZE):
        batch = descriptions_with_codes[batch_idx : batch_idx + BATCH_SIZE]
        logger.info(
            f"🔄 Embedding batch {batch_idx // BATCH_SIZE + 1} "
            f"({len(batch)} texts, total tokens ~{sum(len(t.split()) for t in batch)})"
        )

        try:
            result = client.embed(
                texts=batch,
                model=model,
                input_type="document",  # corpus side of retrieval
            )
        except Exception as e:
            logger.error(f"❌ Voyage API error: {e}")
            sys.exit(1)

        all_embeddings.extend(result.embeddings)
        logger.info(f"  ✅ Batch returned {len(result.embeddings)} embeddings")

    arr = np.array(all_embeddings, dtype=np.float32)
    logger.info(f"📐 Final embeddings shape: {arr.shape}")
    return arr


# ============================================================================
# SAVE
# ============================================================================

def save_embeddings(
    codes: List[str],
    descriptions: List[str],
    categories: List[str],
    embeddings: np.ndarray,
    model: str,
):
    """Persist everything to a single pickle for easy loading by the classifier."""
    payload = {
        "codes": codes,
        "descriptions": descriptions,
        "categories": categories,
        "embeddings": embeddings,
        "metadata": {
            "model": model,
            "dimension": embeddings.shape[1],
            "code_count": len(codes),
            "created_at": datetime.now(timezone.utc).isoformat(),
            "source_file": INPUT_PATH.name,
        },
    }
    with OUTPUT_PATH.open("wb") as f:
        pickle.dump(payload, f, protocol=pickle.HIGHEST_PROTOCOL)
    logger.info(f"💾 Saved to {OUTPUT_PATH} ({OUTPUT_PATH.stat().st_size / 1024:.1f} KB)")


# ============================================================================
# CLI
# ============================================================================

def main():
    import argparse
    parser = argparse.ArgumentParser(description="Embed USPTO design codes via Voyage AI")
    parser.add_argument(
        "--model", default=DEFAULT_MODEL,
        help=f"Voyage model to use (default: {DEFAULT_MODEL}). "
             "Upgrade to voyage-3-large for marginal quality gains."
    )
    args = parser.parse_args()

    codes, descriptions, categories = load_codes()
    texts_to_embed = [
        build_searchable_text(c, d) for c, d in zip(codes, descriptions)
    ]
    embeddings = embed_descriptions(texts_to_embed, model=args.model)

    save_embeddings(codes, descriptions, categories, embeddings, args.model)

    logger.info("\n" + "=" * 60)
    logger.info("📊 EMBEDDING BUILD COMPLETE")
    logger.info("=" * 60)
    logger.info(f"  Model:      {args.model}")
    logger.info(f"  Codes:      {len(codes):,}")
    logger.info(f"  Dimensions: {embeddings.shape[1]}")
    logger.info(f"  Output:     {OUTPUT_PATH}")
    logger.info("=" * 60)
    logger.info("\n💡 Next step: design_code_classifier.py will load this file at startup")


if __name__ == "__main__":
    main()