capture.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 json
import uuid
from datetime import datetime, timezone
from typing import Any
 
import httpx
 
from aquin.types import CaptureBundle
 
 
def build_minimal_capture_bundle(
    prompt: str,
    model: Any,
    session_id: str = "",
    sae_layer: int | None = None,
) -> CaptureBundle:
    """
    Run a forward pass through model and collect top SAE features + logit lens.
    Falls back to empty tensors if SAE weights are not available.
    """
    import torch
    from aquin.compute.model_loader import get_loaded_model, resolve_model_id, get_config
 
    if model is None:
        model = get_loaded_model()
    if model is None:
        raise RuntimeError("No model loaded. Run: aquin load --model <model-id>")
 
    model_id = getattr(model, "cfg", None)
    hf_name = ""
    short = ""
    try:
        hf_name = model.cfg.model_name if hasattr(model, "cfg") else str(type(model).__name__)
        short = resolve_model_id(hf_name)
        cfg = get_config(short)
        if sae_layer is None:
            sae_layer = cfg.get("sae_layer", 0)
        hf_name = cfg["hf_name"]
    except Exception:
        short = "unknown"
        if sae_layer is None:
            sae_layer = 0
 
    # tokenize + forward pass
    tokens = model.to_tokens(prompt, prepend_bos=True)
    with torch.no_grad():
        _, cache = model.run_with_cache(tokens)
 
    # logit lens: top token per layer
    logit_lens: list[dict] = []
    try:
        n_layers = model.cfg.n_layers
        for layer in range(n_layers):
            resid = cache[f"blocks.{layer}.hook_resid_post"][0, -1]
            logits = model.unembed(model.ln_final(resid.unsqueeze(0).unsqueeze(0)))[0, 0]
            top_tok = int(logits.argmax().item())
            top_str = model.to_string([top_tok])
            logit_lens.append({"layer": layer, "token": top_str, "token_id": top_tok})
    except Exception:
        logit_lens = []
 
    # SAE top features
    top_features: list[dict] = []
    try:
        from aquin.compute.model_loader import load_sae
        sae = load_sae(model, sae_layer, short)
        resid = cache[f"blocks.{sae_layer}.hook_resid_post"][0, -1]
        feat_acts, top_idxs = sae.get_top_features(resid, top_k=20)
        top_features = [
            {"feature_idx": int(idx), "activation": float(act)}
            for idx, act in zip(top_idxs.tolist(), feat_acts.tolist())
        ]
    except Exception:
        top_features = []
 
    # greedy decode single token for response
    with torch.no_grad():
        logits = model(tokens)
    next_tok = int(logits[0, -1].argmax().item())
    response = model.to_string([next_tok])
 
    return CaptureBundle(
        version=1,
        capture_id="",
        session_id=session_id,
        model_id=short,
        hf_name=hf_name,
        prompt=prompt,
        response=response,
        created_at=datetime.now(timezone.utc).isoformat(),
        top_features=top_features,
        sae_layer=sae_layer,
        logit_lens=logit_lens,
        attention={},
    )
 
 
def upload_capture(
    bundle: CaptureBundle,
    session_id: str,
    api_key: str,
    base_url: str = "https://api.aquin.app",
) -> str:
    """
    Upload a CaptureBundle to Aquin Cloud. Returns capture_id.
    """
    import os
    base_url = os.environ.get("AQUIN_BASE_URL", base_url).rstrip("/")
 
    payload = bundle.model_copy(update={"session_id": session_id}).model_dump()
    body = json.dumps(payload).encode()
 
    resp = httpx.post(
        f"{base_url}/api/sync/captures/upload",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        },
        content=body,
        timeout=30,
        verify=False,
    )
    resp.raise_for_status()
    capture_id = resp.json()["capture_id"]
    return capture_id