saes.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
from __future__ import annotations
 
import os
import sys
from pathlib import Path
 
 
_BASE_URL = "https://api.aquin.app"
 
 
def _base_url() -> str:
    return os.environ.get("AQUIN_BASE_URL", _BASE_URL).rstrip("/")
 
 
def _get_api_key(cfg: dict) -> str:
    key = cfg.get("api_key") or os.environ.get("AQUIN_API_KEY") or ""
    if not key:
        print("No API key found. Set AQUIN_API_KEY, or use aquin load sae --path <file.pt>.")
        sys.exit(1)
    return key
 
 
# ── Public (no auth) ─────────────────────────────────────────────────────────
 
def list_saes(model: str | None = None) -> list[dict]:
    import requests
    from aquin.user_errors import die
 
    try:
        resp = requests.get(f"{_base_url()}/api/public/saes", timeout=15)
        resp.raise_for_status()
    except Exception as exc:
        die(message="Could not load SAE catalog. Try again shortly.", exc=exc, label="load sae")
    data = resp.json()
    if model:
        data = [s for s in data if s["model_slug"] == model or s["model_name"].lower() == model.lower()]
    return data
 
 
def get_sae(id: str) -> dict:
    import requests
    from aquin.user_errors import die, sanitize_error_string
 
    resp = requests.get(f"{_base_url()}/api/public/saes/{id}", timeout=15)
    if resp.status_code == 404:
        die(sanitize_error_string(f"SAE '{id}' not found"))
    try:
        resp.raise_for_status()
    except Exception as exc:
        die(message="Could not load SAE info. Try again shortly.", exc=exc, label="info sae")
    return resp.json()
 
 
# ── Authenticated ─────────────────────────────────────────────────────────────
 
def pull_sae(id: str, output: str | Path, api_key: str) -> None:
    import requests
 
    out = Path(output)
    out.parent.mkdir(parents=True, exist_ok=True)
 
    resp = requests.get(
        f"{_base_url()}/api/public/saes/{id}/pull",
        headers={"Authorization": f"Bearer {api_key}"},
        stream=True,
        timeout=600,
        allow_redirects=True,
    )
 
    if resp.status_code == 503:
        out.unlink(missing_ok=True)
        print(f"Error: SAE '{id}' is invalid in catalog storage (not real weights).", file=sys.stderr)
        print("Contact Aquin support or use: aquin load sae --path <local.pt> --layer <n>", file=sys.stderr)
        sys.exit(1)
    if resp.status_code == 401:
        from aquin.auth_config import print_token_rejected
        print_token_rejected(api_key)
    if resp.status_code == 403:
        data = resp.json()
        print(data.get("error", "Access not granted."))
        sys.exit(1)
    if resp.status_code == 404:
        print(f"SAE '{id}' is not in the Aquin SAE library.", file=sys.stderr)
        print("", file=sys.stderr)
        print("Try:  aquin load sae <model-l{n}>", file=sys.stderr)
        print("      aquin load sae --path <file.pt> --layer <n>", file=sys.stderr)
        print("", file=sys.stderr)
        print(
            "Train your own: aquin sae train --layer <n> --name <tag>",
            file=sys.stderr,
        )
        sys.exit(1)
    try:
        resp.raise_for_status()
    except Exception as exc:
        from aquin.user_errors import die
        die(
            message=f"Could not download SAE '{id}'. Set AQUIN_API_KEY or use --path <file.pt>.",
            exc=exc,
            label="load sae",
        )
 
    total = int(resp.headers.get("content-length", 0))
    downloaded = 0
    with open(out, "wb") as f:
        for chunk in resp.iter_content(chunk_size=1024 * 1024):
            f.write(chunk)
            downloaded += len(chunk)
            if total:
                pct = downloaded / total * 100
                mb = downloaded / 1_000_000
                print(f"\r  {mb:.1f} MB / {total/1_000_000:.1f} MB  ({pct:.0f}%)", end="", flush=True)
    print(f"\r  Saved to {out}                          ")
 
    from aquin.compute.torch_io import (
        is_valid_sae_checkpoint_path,
        looks_like_invalid_norm_blob,
        looks_like_json_error_blob,
    )
 
    if looks_like_json_error_blob(out) or not is_valid_sae_checkpoint_path(out):
        out.unlink(missing_ok=True)
        print(f"Error: download for '{id}' was not a valid SAE checkpoint.", file=sys.stderr)
        print("Catalog storage may be corrupt — train locally:", file=sys.stderr)
        print("  aquin sae train --model <id> --layer <n> --name <tag> --quick", file=sys.stderr)
        print("  aquin load sae --user <tag> --layer <n>", file=sys.stderr)
        sys.exit(1)
 
    try:
        row = get_sae(id)
    except Exception:
        row = {}
    if row.get("meta_key"):
        layer = row.get("layer")
        norm_name = f"norm_layer{layer}.pt" if layer is not None else "norm.pt"
        norm_out = out.parent / norm_name
        norm_resp = requests.get(
            f"{_base_url()}/api/public/saes/{id}/pull-norm",
            headers={"Authorization": f"Bearer {api_key}"},
            stream=True,
            timeout=600,
            allow_redirects=True,
        )
        if norm_resp.status_code == 200:
            with open(norm_out, "wb") as f:
                for chunk in norm_resp.iter_content(chunk_size=1024 * 1024):
                    f.write(chunk)
            if looks_like_invalid_norm_blob(norm_out):
                norm_out.unlink(missing_ok=True)
                print(
                    f"  Note: norm stats unavailable for '{id}' (catalog storage issue) — inspect works without them",
                    file=sys.stderr,
                )
            else:
                print(f"  Saved norm to {norm_out}")
        elif norm_resp.status_code == 503:
            print(
                f"  Warning: norm for '{id}' is invalid in catalog storage — inspect will run without normalization",
                file=sys.stderr,
            )
        elif norm_resp.status_code != 404:
            print(f"  Warning: norm download failed ({norm_resp.status_code})", file=sys.stderr)
 
 
# ── SDK client namespace ──────────────────────────────────────────────────────
 
class SaesNamespace:
    def __init__(self, api_key: str | None = None):
        self._api_key = api_key
 
    def list(self, model: str | None = None) -> list[dict]:
        return list_saes(model=model)
 
    def get(self, id: str) -> dict:
        return get_sae(id)
 
    def pull(self, id: str | None = None, model: str | None = None,
             layer: int | None = None, output: str | Path = "./sae.pt") -> None:
        if id is None:
            if model is None:
                raise ValueError("Provide id or model.")
            rows = list_saes(model=model)
            if layer is not None:
                rows = [r for r in rows if r["layer"] == layer]
            if not rows:
                raise ValueError(f"No SAE found for model={model} layer={layer}")
            id = rows[0]["id"]
        if not self._api_key:
            raise ValueError("api_key required for pull(). Pass it to aquin.Client().")
        pull_sae(id, output, self._api_key)