""" 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()