data.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.
 
"""Data — activation capture / replay, feature compare."""
 
from __future__ import annotations
 
import json
from pathlib import Path
from typing import Any
 
from aquin.sdk._runtime import invoke_direct, prepare_compute
 
 
def _load_probes(probes: Any) -> list[dict[str, Any]]:
    if probes is None:
        return []
    if isinstance(probes, list):
        return list(probes)
    path = Path(str(probes)).expanduser()
    text = path.read_text(encoding="utf-8")
    if path.suffix == ".jsonl":
        rows = []
        for line in text.splitlines():
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            rows.append(row if isinstance(row, dict) else {"text": str(row)})
        return rows
    data = json.loads(text)
    if isinstance(data, list):
        return [r if isinstance(r, dict) else {"text": str(r)} for r in data]
    if isinstance(data, dict) and "probes" in data:
        return list(data["probes"])
    raise ValueError(f"Unsupported probes file: {path}")
 
 
def capture_activations(
    *,
    output_dir: str | Path,
    probes: Any = None,
    model_id: str | None = None,
    **kwargs: Any,
) -> dict[str, Any]:
    """Record activations (+ metadata) into ``output_dir``."""
    from aquin.compute.activation_capture import run_capture_activations
    from aquin.compute.model_loader import get_active_model_id
 
    mid = model_id or get_active_model_id()
    if not mid:
        raise ValueError("No model loaded. Call aquin.sdk.session.load_model(...) first.")
    prepare_compute(model_id=mid)
    probe_rows = _load_probes(probes)
    if not probe_rows:
        raise ValueError("probes= is required (list of dicts or path to json/jsonl).")
    return invoke_direct(
        run_capture_activations,
        mid,
        probe_rows,
        output_dir,
        command="activations capture",
        tool_name="run_capture_activations",
        **kwargs,
    )
 
 
def replay_activations(
    path: str | Path,
    *,
    filter_expr: str | None = None,
    group_by: str | None = None,
    limit: int | None = None,
) -> dict[str, Any]:
    """Slice a saved capture by metadata without re-running the model."""
    from aquin.compute.activation_replay import replay_capture
 
    return invoke_direct(
        replay_capture,
        path,
        command="activations replay",
        tool_name="activations_replay",
        filter_expr=filter_expr,
        group_by=group_by,
        limit=limit,
    )
 
 
def features_compare(**kwargs: Any) -> dict[str, Any]:
    from aquin.features_cli import cmd_features
 
    argv: list[str] = ["compare"]
    for key, flag in (
        ("from_dir", "--from"),
        ("source", "--from"),
        ("group", "--group"),
        ("encode_sae", "--encode-sae"),
    ):
        if kwargs.get(key) is not None:
            argv.extend([flag, str(kwargs[key])])
    cmd_features(argv)
    return {"ok": True}