loader_shim.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.
 
"""
Patches ingested compute modules so load_model/load_sae delegate to
aquin.compute.model_loader — same model object as aquin load.
Call apply() once at engine startup before any tool invocation.
"""
from __future__ import annotations
 
_applied = False
 
 
def apply() -> None:
    global _applied
    if _applied:
        return
    _applied = True
 
    import aquin.compute.causal_trace as _ct
    import aquin.compute.feature_analysis as _fa
    import aquin.compute.model_loader as _ml
 
    # Patch causal_trace.load_model → model_loader.load_model (same signature)
    _ct.load_model = _ml.load_model
 
    # Patch feature_analysis.load_sae to accept (model_id, layer=None) → _ml.load_sae
    def _shim_load_sae(model_id: str, layer: int | None = None):
        short = _ml.resolve_model_id(model_id)
        cfg = _ml.get_config(short)
        resolved_layer = int(layer if layer is not None else cfg["sae_layer"])
        if layer is not None:
            resolved_layer = _ml.require_sae_layer(
                short,
                resolved_layer,
                command="trace",
            )
        model = _ml.get_loaded_model()
        if model is None:
            model = _ml.load_model(short)
        return _ml.load_sae(model, resolved_layer, short)
 
    _fa.load_sae = _shim_load_sae