pipelines.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.
 
"""
Multi-step orchestrations that combine multiple compute calls.
No new algorithms — these are the same sub-call sequences as the legacy
runFullPipeline flow, exposed as importable functions so bridge.py and raw
CLI commands can call them without the tool registry.
"""
from __future__ import annotations
 
from typing import Any
 
 
def run_full_inspection(
    prompt: str,
    model_id: str,
    ctx: dict,
    *,
    layer: int | None = None,
) -> dict:
    """
    chat + feature_analysis + logit_lens — same sub-calls as legacy runFullPipeline.
    Returns a dict with keys: prompt, response, model_id, top_features, sae_layer,
    logit_lens, trace_results, color_map, prompt_tokens, response_tokens, attribution.
    """
    from aquin.compute.model_loader import (
        get_loaded_model,
        load_model,
        require_sae_layer,
        resolve_model_id,
    )
    from aquin.compute.causal_trace import run_chat, run_logit_lens
    from aquin.compute.feature_analysis import run_feature_analysis_unlabeled
 
    try:
        short = resolve_model_id(model_id)
    except Exception as e:
        return {"error": f"Unknown model '{model_id}': {e}"}
 
    try:
        sae_layer = require_sae_layer(
            short,
            layer,
            command="trace",
            example_suffix=f'--prompt "{prompt[:60]}"' if prompt else "",
        )
    except ValueError as e:
        return {"error": str(e)}
 
    model = get_loaded_model()
    if model is None:
        try:
            model = load_model(short)
        except Exception as e:
            return {"error": f"Could not load model: {e}"}
 
    # 1. Generate response
    try:
        response = run_chat(prompt, model_id=short, max_new_tokens=200, temperature=0.7)
    except Exception as e:
        return {"error": f"Generation failed: {e}"}
 
    # 2. SAE feature analysis
    try:
        feat_result = run_feature_analysis_unlabeled(
            prompt, response, model, model_id=short, layer=sae_layer,
        )
    except Exception as e:
        return {"error": f"Feature analysis failed: {e}"}
 
    actual_layer = feat_result.get("sae_layer")
    if actual_layer is not None and int(actual_layer) != int(sae_layer):
        return {
            "error": (
                f"SAE layer mismatch: requested layer {sae_layer}, "
                f"feature analysis ran at layer {actual_layer}."
            ),
        }
 
    prompt_tokens: list[str] = feat_result.get("prompt_tokens", [])
    response_tokens: list[str] = feat_result.get("response_tokens", [])
    top_features: list[dict] = feat_result.get("top_response_features", [])
    attribution: list[dict] = feat_result.get("attribution", [])
    sae_layer = int(feat_result.get("sae_layer", sae_layer))
 
    # 3. Logit lens
    try:
        logit_lens = run_logit_lens(prompt, model_id=short, top_k=5)
    except Exception as e:
        print(f"[pipelines] logit-lens failed: {e}", flush=True)
        logit_lens = []
 
    # Build trace results (probability delta per layer)
    trace_results: list[dict] = []
    if logit_lens:
        final_prob = logit_lens[-1]["top_tokens"][0]["prob"] if logit_lens[-1].get("top_tokens") else 0.0
        for i, row in enumerate(logit_lens):
            prob = row["top_tokens"][0]["prob"] if row.get("top_tokens") else 0.0
            next_row = logit_lens[i + 1] if i + 1 < len(logit_lens) else None
            next_prob = next_row["top_tokens"][0]["prob"] if next_row and next_row.get("top_tokens") else prob
            delta = max(0.0, next_prob - prob)
            trace_results.append({
                "layer": row["layer"],
                "drop": round(delta, 4),
                "attn_drop": round(delta * 0.6, 4),
                "mlp_drop": round(delta * 0.4, 4),
                "baseline_prob": round(final_prob, 4),
            })
 
    # Build color map from attribution
    color_map: dict | None = None
    feature_color: dict[int, int] = {}
    counter = [0]
 
    def _color(fidx: int) -> int:
        if fidx not in feature_color:
            feature_color[fidx] = counter[0] % 10
            counter[0] += 1
        return feature_color[fidx]
 
    prompt_cm: dict[int, int] = {}
    response_cm: dict[int, int] = {}
    for entry in attribution:
        ri = entry["response_ti"]
        for feat in entry.get("driven_by_features", []):
            cidx = _color(feat["feature_idx"])
            response_cm[ri] = cidx
            for pi in feat.get("also_in_prompt_positions", []):
                prompt_cm[pi] = cidx
 
    if prompt_cm or response_cm:
        color_map = {"prompt": prompt_cm, "response": response_cm}
 
    trace_target = response_tokens[0].strip() if response_tokens else ""
 
    return {
        "prompt": prompt,
        "response": response,
        "model_id": short,
        "top_features": top_features,
        "prompt_tokens": prompt_tokens,
        "response_tokens": response_tokens,
        "attribution": attribution,
        "sae_layer": sae_layer,
        "logit_lens": logit_lens,
        "trace_results": trace_results,
        "trace_target": trace_target,
        "color_map": color_map,
    }