"""
Original from: https://github.com/jonhanke/nam_causal-head-gating
Adapted and extended for Causal Summer project by:
- Bowen
- 2026-03-09

`inteprelens` is the extracted core package for transformer architecture analysis.

The current public surface keeps the CHG APIs intact while the package is being
split out of the larger experiment repository.

Quick start:
    >>> from inteprelens import LensAPI
    >>> analyzer = LensAPI.from_pretrained("meta-llama/Llama-3.2-1B")
    >>> results = analyzer.fit(texts=prompts, targets=targets)
    >>> print(results.necessary_heads())
"""

from .api import LensAPI
from .core.attribution import CausalCircuitAttribution
from .core.chg import CHG
from .core.hooks import (
    autograd_node_hook,
    capture_backward_pre,
    capture_gradients,
    capture_kwargs_hook,
    tensor_grad_hook,
)
from .core.interventions import InterventionSpec, apply_interventions
from .core.probabilities import (
    apply_temperature,
    gather_token_scores,
    logits_to_log_probs,
    logits_to_probs,
    select_final_token,
)
from .core.sites import ALL_SITES, GLOBAL_SITES, LAYER_SITES, SITE_SPECS
from .core.tracer import TransformerTracer
from .core.trainer import CHGTrainer
from .core.types import LayerSiteKey, SiteRequest, TokenScores, TraceResult
from .taxonomy.masks import CHGResults
from .utils.ds import CHGDataset, MaskedSequenceDataset
from .utils.sentence_spans import aggregate_to_sentences, get_sentence_token_spans

__version__ = "0.1.0"
__author__ = "Bowen"
__all__ = [
    "LensAPI",
    "CHG",
    "CHGTrainer",
    "TransformerTracer",
    "TraceResult",
    "SiteRequest",
    "LayerSiteKey",
    "ALL_SITES",
    "LAYER_SITES",
    "GLOBAL_SITES",
    "SITE_SPECS",
    "InterventionSpec",
    "apply_interventions",
    "TokenScores",
    "logits_to_log_probs",
    "logits_to_probs",
    "apply_temperature",
    "select_final_token",
    "gather_token_scores",
    "CausalCircuitAttribution",
    "autograd_node_hook",
    "capture_backward_pre",
    "capture_gradients",
    "capture_kwargs_hook",
    "tensor_grad_hook",
    "CHGDataset",
    "MaskedSequenceDataset",
    "CHGResults",
    "get_sentence_token_spans",
    "aggregate_to_sentences",
]
