publish_public_sae.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Upload SAE + norm + UMAP to R2 and register in Supabase public_saes (no Node required)."""
 
from __future__ import annotations
 
import json
import os
import sys
from pathlib import Path
from typing import Any
 
 
def _load_env_files() -> None:
    roots = [
        Path.cwd(),
        Path.cwd().parent,
        Path(__file__).resolve().parents[2],
    ]
    seen: set[Path] = set()
    for root in roots:
        for name in (".env.local", ".env"):
            path = (root / name).resolve()
            if path in seen or not path.is_file():
                continue
            seen.add(path)
            for line in path.read_text(encoding="utf-8").splitlines():
                line = line.strip()
                if not line or line.startswith("#") or "=" not in line:
                    continue
                key, val = line.split("=", 1)
                key = key.strip()
                val = val.strip().strip('"').strip("'")
                os.environ.setdefault(key, val)
 
 
def _parse_flag(args: list[str], name: str) -> str | None:
    for i, a in enumerate(args):
        if a == name and i + 1 < len(args):
            return args[i + 1]
    return None
 
 
def _parse_num(args: list[str], name: str, fallback: float | int | None = None) -> float | int | None:
    raw = _parse_flag(args, name)
    if raw is None or raw == "":
        return fallback
    try:
        n = float(raw)
        return int(n) if n == int(n) else n
    except ValueError:
        return fallback
 
 
def _read_json(path: Path) -> dict[str, Any] | None:
    if not path.is_file():
        return None
    try:
        return json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return None
 
 
def _r2_client():
    try:
        import boto3
    except ImportError as exc:
        raise RuntimeError("pip install boto3") from exc
 
    endpoint = os.environ.get("R2_ENDPOINT")
    access_key = os.environ.get("R2_ACCESS_KEY_ID")
    secret = os.environ.get("R2_SECRET_ACCESS_KEY")
    missing = [
        name
        for name, val in (
            ("R2_ENDPOINT", endpoint),
            ("R2_ACCESS_KEY_ID", access_key),
            ("R2_SECRET_ACCESS_KEY", secret),
            ("R2_BUCKET_NAME", os.environ.get("R2_BUCKET_NAME")),
        )
        if not val
    ]
    if missing:
        raise RuntimeError(f"Missing env: {', '.join(missing)} (set in .env.local)")
 
    return boto3.client(
        "s3",
        endpoint_url=endpoint,
        aws_access_key_id=access_key,
        aws_secret_access_key=secret,
        region_name="auto",
    )
 
 
def _upload_r2(client: Any, bucket: str, key: str, file_path: Path, content_type: str) -> None:
    body = file_path.read_bytes()
    client.put_object(Bucket=bucket, Key=key, Body=body, ContentType=content_type)
    print(f"  uploaded {key} ({len(body) / 1_000_000:.2f} MB)", flush=True)
 
 
def _supabase_upsert(row: dict[str, Any]) -> dict[str, Any]:
    try:
        import requests
    except ImportError as exc:
        raise RuntimeError("pip install requests") from exc
 
    url = os.environ.get("NEXT_PUBLIC_SUPABASE_URL") or os.environ.get("SUPABASE_URL")
    key = os.environ.get("SUPABASE_SERVICE_ROLE_KEY")
    if not url or not key:
        raise RuntimeError("Set NEXT_PUBLIC_SUPABASE_URL and SUPABASE_SERVICE_ROLE_KEY in .env.local")
 
    base = url.rstrip("/")
    headers = {
        "apikey": key,
        "Authorization": f"Bearer {key}",
        "Content-Type": "application/json",
        "Prefer": "resolution=merge-duplicates,return=representation",
    }
    resp = requests.post(
        f"{base}/rest/v1/public_saes?on_conflict=model_slug,layer",
        headers=headers,
        json=row,
        timeout=60,
    )
    if resp.status_code in (200, 201):
        data = resp.json()
        return data[0] if isinstance(data, list) and data else data
 
    # Fallback: plain insert
    resp2 = requests.post(
        f"{base}/rest/v1/public_saes",
        headers={**headers, "Prefer": "return=representation"},
        json=row,
        timeout=60,
    )
    if resp2.status_code not in (200, 201):
        raise RuntimeError(f"Supabase error {resp.status_code}: {resp.text}; insert: {resp2.text}")
    data = resp2.json()
    return data[0] if isinstance(data, list) and data else data
 
 
def publish_public_sae(args: list[str]) -> None:
    _load_env_files()
 
    sae_path_s = _parse_flag(args, "--sae")
    if not sae_path_s:
        print(
            "Usage: aquin sae publish --sae <path.pt> --model-slug <slug> "
            "[--layer N] [--model-name ...] [--d-model 1024] [--d-sae 32768]",
            file=sys.stderr,
        )
        sys.exit(1)
 
    sae_abs = Path(sae_path_s).expanduser().resolve()
    if not sae_abs.is_file():
        print(f"SAE not found: {sae_abs}", file=sys.stderr)
        sys.exit(1)
 
    meta = _read_json(sae_abs.with_suffix(".meta.json")) or {}
    model_slug = _parse_flag(args, "--model-slug") or str(meta.get("model_id") or "").strip()
    if not model_slug:
        print("Pass --model-slug or ensure .meta.json has model_id", file=sys.stderr)
        sys.exit(1)
 
    layer_raw = _parse_num(args, "--layer") or meta.get("layer")
    if layer_raw is None:
        print("Pass --layer or ensure .meta.json has layer", file=sys.stderr)
        sys.exit(1)
    layer = int(layer_raw)
 
    d_model = int(_parse_num(args, "--d-model", meta.get("d_model") or 2048) or 2048)
    d_sae = int(_parse_num(args, "--d-sae", meta.get("n_features") or 32768) or 32768)
    train_steps = _parse_num(args, "--train-steps", meta.get("steps"))
    l0 = _parse_num(args, "--l0")
    dead_pct = _parse_num(args, "--dead-pct")
    interp_score = _parse_num(args, "--interp-score")
    model_name = _parse_flag(args, "--model-name") or model_slug
    arch = _parse_flag(args, "--arch") or "sparse_autoencoder"
    k = int(_parse_num(args, "--k") or round(d_sae / d_model))
 
    metrics = _read_json(sae_abs.with_suffix(".metrics.json"))
    if metrics and interp_score is None and metrics.get("interp_score") is not None:
        interp_score = float(metrics["interp_score"])
        print(f"  metrics sidecar: {sae_abs.with_suffix('.metrics.json').name}", flush=True)
 
    umap_path_s = _parse_flag(args, "--umap")
    umap_path = Path(umap_path_s).expanduser() if umap_path_s else sae_abs.with_suffix(".umap.json")
    if not umap_path.is_file():
        umap_path = None  # type: ignore[assignment]
 
    norm_path_s = _parse_flag(args, "--norm")
    if norm_path_s:
        norm_path = Path(norm_path_s).expanduser()
    else:
        parent = sae_abs.parent
        candidates = [
            parent / f"norm_layer{layer}.pt",
            parent / f"_acts_layer{layer}" / "norm.pt",
            parent / "norm.pt",
        ]
        norm_path = next((p for p in candidates if p.is_file()), None)
    if norm_path is None or not norm_path.is_file():
        print("Norm file not found. Pass --norm or train with aquin sae train.", file=sys.stderr)
        sys.exit(1)
 
    bucket = os.environ["R2_BUCKET_NAME"]
    r2_sae_key = f"saes/{model_slug}/l{layer}/sae.pt"
    r2_norm_key = f"saes/{model_slug}/l{layer}/norm.pt"
    r2_umap_key = f"saes/{model_slug}/l{layer}/umap.json"
 
    print(f"Publishing {model_slug} layer {layer}...", flush=True)
    print("R2:", flush=True)
    s3 = _r2_client()
    _upload_r2(s3, bucket, r2_sae_key, sae_abs, "application/octet-stream")
    _upload_r2(s3, bucket, r2_norm_key, norm_path.resolve(), "application/octet-stream")
    umap_key = _parse_flag(args, "--umap-key")
    if umap_path and umap_path.is_file():
        _upload_r2(s3, bucket, r2_umap_key, umap_path, "application/json")
        umap_key = umap_key or r2_umap_key
        print(f"  umap: {umap_path.name}", flush=True)
    else:
        print("  umap: (skipped — run aquin sae catalog-metrics first)", flush=True)
 
    row = {
        "model_slug": model_slug,
        "model_name": model_name,
        "layer": layer,
        "d_model": d_model,
        "d_sae": d_sae,
        "k": k,
        "arch": arch,
        "train_steps": train_steps,
        "final_loss": _parse_num(args, "--final-loss"),
        "l0": l0,
        "dead_pct": dead_pct,
        "interp_score": interp_score,
        "umap_key": umap_key,
        "r2_key": r2_sae_key,
        "meta_key": r2_norm_key,
    }
 
    data = _supabase_upsert(row)
    print(f"\nRegistered: {data}", flush=True)
    print(f"Pull: aquin load sae {model_slug}-l{layer}", flush=True)
    print(f"Load: aquin load sae {model_slug}-l{{n}}", flush=True)