`visual-search/embed_shopify_gap.py`

115 lines
#!/usr/bin/env python3
# Embed the ACTIVE Shopify products that have no CLIP vector yet ("the 46% gap").
# Sibling of embed_catalog.py — same model/weights/normalization so vectors share the
# catalog's space — but sources from shopify_products (Shopify CDN image_url) instead of
# vendor_catalog, because ~97.5% of the gap SKUs have no vendor_catalog image row.
#
# Key strategy: image_embeddings PK is vc_id (vendor_catalog.id). Shopify-only rows use a
# SYNTHETIC NEGATIVE key vc_id = -shopify_id — bigint, can never collide with real
# vendor_catalog ids, no schema change. search_service.load() picks rows up unchanged
# (it selects by `embedding is not null` and indexes by dw_sku).
#
# Resumable: a dw_sku with any embedding-NOT-NULL row is skipped; permanent failures get
# the NULL-embedding placeholder on the SAME synthetic key so reruns skip them too.
#
# Usage:  DW_UNIFIED_DB=<uri> python3 embed_shopify_gap.py [--limit N] [--batch 32] [--workers 24]
import os, sys, io, time, argparse, concurrent.futures as cf
import urllib.request
import numpy as np, psycopg2
from PIL import Image
import torch, open_clip

DB = os.environ.get("DW_UNIFIED_DB") or sys.exit("set DW_UNIFIED_DB")
ap = argparse.ArgumentParser()
ap.add_argument("--limit", type=int, default=0)     # 0 = all remaining
ap.add_argument("--batch", type=int, default=32)    # CLIP forward-pass batch
ap.add_argument("--workers", type=int, default=24)  # concurrent downloads
ap.add_argument("--chunk", type=int, default=2000)  # rows fetched per DB round
args = ap.parse_args()

print("loading CLIP ViT-B-32 (laion2b)…", flush=True)
model, _, preprocess = open_clip.create_model_and_transforms("ViT-B-32", pretrained="laion2b_s34b_b79k")

_ft = os.path.join(os.path.dirname(os.path.abspath(__file__)), "dw_clip_ft.pt")
if os.path.exists(_ft):
    model.load_state_dict(torch.load(_ft, map_location="cpu"), strict=False)
    print("loaded fine-tuned DW CLIP weights:", _ft, flush=True)

model.eval()
torch.set_num_threads(max(1, os.cpu_count() - 2))
DIM = model.visual.output_dim

def fetch(row):
    shopify_id, dw_sku, vendor, pattern, url = row
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "dw-visual/1.0"})
        with urllib.request.urlopen(req, timeout=8) as r:
            data = r.read(6_000_000)
        img = Image.open(io.BytesIO(data)).convert("RGB")
        return (shopify_id, dw_sku, vendor, pattern, url, preprocess(img))
    except Exception:
        return (shopify_id, dw_sku, vendor, pattern, url, None)   # permanent-fail placeholder

# One row per gap dw_sku: the newest ACTIVE shopify product carrying that sku.
# shopify_id is a GID string ("gid://shopify/Product/779…") — extract the numeric
# tail here so the synthetic key (-id) is a clean bigint.
GAP_SQL = """
WITH act AS (
  SELECT DISTINCT ON (sp.dw_sku)
         (regexp_replace(sp.shopify_id, '\\D', '', 'g'))::bigint AS shopify_id,
         sp.dw_sku, sp.vendor, sp.pattern_name, sp.image_url
  FROM shopify_products sp
  WHERE sp.status='ACTIVE' AND sp.image_url IS NOT NULL AND sp.handle IS NOT NULL AND sp.dw_sku IS NOT NULL
    AND sp.shopify_id ~ '\\d'
  ORDER BY sp.dw_sku, sp.synced_at DESC NULLS LAST
),
emb AS (SELECT DISTINCT dw_sku FROM image_embeddings WHERE embedding IS NOT NULL),
tried AS (SELECT DISTINCT dw_sku FROM image_embeddings WHERE vc_id < 0 AND embedding IS NULL)
SELECT a.shopify_id, a.dw_sku, a.vendor, a.pattern_name, a.image_url
FROM act a
LEFT JOIN emb e ON e.dw_sku = a.dw_sku
LEFT JOIN tried t ON t.dw_sku = a.dw_sku
WHERE e.dw_sku IS NULL AND t.dw_sku IS NULL
LIMIT %s
"""

def main():
    conn = psycopg2.connect(DB); conn.autocommit = False
    total_done = 0; t0 = time.time()
    while True:
        cur = conn.cursor()
        cur.execute(GAP_SQL, (args.chunk,))
        rows = cur.fetchall(); cur.close()
        if not rows: break
        with cf.ThreadPoolExecutor(max_workers=args.workers) as ex:
            got = list(ex.map(fetch, rows))
        ok  = [g for g in got if g[5] is not None]
        bad = [g for g in got if g[5] is None]
        recs = []
        for i in range(0, len(ok), args.batch):
            b = ok[i:i+args.batch]
            with torch.no_grad():
                t = torch.stack([g[5] for g in b])
                v = model.encode_image(t)
                v = (v / v.norm(dim=-1, keepdim=True)).cpu().numpy().astype("float32")
            for g, emb in zip(b, v):
                # vc_id = -shopify_id  (synthetic key, see header)
                recs.append((-g[0], g[1], None, g[2], g[3], g[4], DIM, psycopg2.Binary(emb.tobytes())))
        cur = conn.cursor()
        if recs:
            cur.executemany("""insert into image_embeddings (vc_id,dw_sku,mfr_sku,vendor_code,pattern_name,image_url,dim,embedding)
                               values (%s,%s,%s,%s,%s,%s,%s,%s) on conflict (vc_id) do nothing""", recs)
        if bad:
            cur.executemany("""insert into image_embeddings (vc_id,dw_sku,image_url,embedding) values (%s,%s,%s,NULL)
                               on conflict (vc_id) do nothing""", [(-g[0], g[1], g[4]) for g in bad])
        conn.commit(); cur.close()
        total_done += len(rows)
        rate = total_done / max(1e-6, time.time()-t0)
        print(f"  +{len(rows)} (ok={len(ok)} fail={len(bad)}) · {total_done} done · {rate:.1f}/s", flush=True)
        if args.limit and total_done >= args.limit: break
    conn.close()
    print(f"done: {total_done} processed in {time.time()-t0:.0f}s", flush=True)

if __name__ == "__main__":
    main()