Utils API¶

`deeplens.utils.analysis` ¶

`AnalysisUtils` ¶

Unified analysis utilities for model logits and SAE features.

This class provides methods for visualizing and analyzing transformer model outputs, including logit heatmaps, token probability distributions, and sparse autoencoder (SAE) feature extraction.

Source code in deeplens\utils\analysis.py

class AnalysisUtils():
    """Unified analysis utilities for model logits and SAE features.

    This class provides methods for visualizing and analyzing transformer model
    outputs, including logit heatmaps, token probability distributions, and
    sparse autoencoder (SAE) feature extraction.
    """
    def __init__(
            self,
            hf_model: str = None,
            sae_model: str = None,
            sae_config: str | dict = None,
            layer: int = None,
            cache_dir: str = 'cache'
        ):
        """Initialize the AnalysisUtils class for analyzing sparse autoencoder results.

        Args:
            hf_model (str, optional): Name or path of the HuggingFace model for extracting
                MLP activations (e.g., "gpt2", "meta-llama/Llama-2-7b"). Required for
                `get_most_active_features`. Defaults to None.
            sae_model (str, optional): Path to the trained sparse autoencoder model weights
                file. Required for `get_most_active_features`. Defaults to None.
            sae_config (str | dict, optional): Path to the YAML configuration file or a
                dictionary containing SAE hyperparameters. Required for 
                `get_most_active_features`. Defaults to None.
            layer (int, optional): Index of the transformer layer to extract activations
                from (0-indexed). Required for `get_most_active_features`. Defaults to None.
            cache_dir (str, optional): Directory to cache downloaded models.
                Defaults to 'cache'.
        """
        self.tokenizer = AutoTokenizer.from_pretrained(hf_model, cache_dir=cache_dir)
        self.hf_model = hf_model
        self.sae_model = sae_model
        self.sae_config = sae_config
        self.layer = layer

        self._mlp_extractor = None
        self._sae_extractor = None

    @property
    def mlp_extractor(self):
        if self._mlp_extractor is None:
            if self.hf_model is None or self.layer is None:
                raise ValueError("hf_model and layer required for MLP extraction")
            self._mlp_extractor = ExtractSingleSample(hf_model=self.hf_model, layer=self.layer)
        return self._mlp_extractor

    @property
    def sae_extractor(self):
        if self._sae_extractor is None:
            if self.sae_model is None or self.sae_config is None:
                raise ValueError("sae_model and sae_config required for SAE extraction")
            self._sae_extractor = InterveneFeatures(
                sae_model=self.sae_model,
                sae_config=self.sae_config
            )
        return self._sae_extractor

    def generate_logit_heatmap(
            self,
            logits: torch.Tensor, 
            save_name: str = None,
            k: int | None = None
        ) -> None:
        """Generate and display a heatmap visualization of logits across token positions.

        Creates a color-coded heatmap showing the distribution of logit values across the
        vocabulary for each token position in a sequence. Uses the 'inferno' colormap for
        visualization.

        Args:
            logits (torch.Tensor): Logits tensor with shape (sequence_length, vocab_size)
                or similar. Will be automatically squeezed and moved to CPU if needed.
            save_name (str, optional): Filename (without extension) to save the plot.
                If provided, saves the figure as a PNG file with 300 DPI. If None, only
                displays the plot. Defaults to None.
            k (int): top k vocabulary indexes to plot. Defaults to None.

        Returns:
            None: Displays the plot and optionally saves it to disk.

        Note:
            The figure size is set to (20, 6) for optimal visualization of typical
            sequence lengths and vocabulary sizes.
        """
        if isinstance(logits, torch.Tensor):
            logits = logits.squeeze().detach().cpu().numpy()

        if k is not None:
            if isinstance(logits, np.ndarray):
                logits = torch.from_numpy(logits)
            logits = torch.topk(logits, k=k, dim=-1).values.numpy()

        plt.figure(figsize=(20, 6))
        plt.imshow(logits, cmap="inferno", aspect="auto")
        plt.colorbar()
        plt.xlabel("Vocabulary Index", size=18)
        plt.ylabel("Token Position", size=18)
        plt.xticks(size=12)
        plt.yticks(size=12)
        if save_name is not None:
            plt.savefig(f"{save_name}.png", dpi=300, bbox_inches="tight")
        plt.show()

    def plot_topk_distribution(
            self,
            logits: torch.Tensor, 
            k: int = 20, 
            position: int = 0, 
            save_name: str = None,
            use_softmax: bool = True,
            title: str = None
        ) -> None:
        """Plot a bar chart showing the top-k most probable tokens at a specific sequence position.

        Creates a horizontal bar chart displaying the highest-probability token predictions
        at a given position in the sequence. Useful for analyzing model predictions and
        understanding which tokens are most likely at specific positions.

        Args:
            logits (torch.Tensor): Logits tensor with shape (sequence_length, vocab_size)
                or similar. Will be automatically squeezed and moved to CPU if needed.
            k (int, optional): Number of top predictions to display in the bar chart.
                Defaults to 20.
            position (int, optional): Token position in the sequence to analyze. Use 0-based
                indexing. Defaults to 0 (first token).
            save_name (str, optional): Filename (without extension) to save the plot.
                If provided, saves the figure as a PNG file with 300 DPI. If None, only
                displays the plot. Defaults to None.
            use_softmax (bool, optional): If True, applies softmax to convert logits to
                probabilities before plotting. If False, plots raw logit values.
                Defaults to True.
            title (str, optional): Custom title for the plot. If None, no title is displayed.
                Defaults to None.

        Returns:
            None: Displays the plot and optionally saves it to disk.

        Note:
            Token labels are rotated 45 degrees for readability. The y-axis label changes
            based on use_softmax: shows probability notation if True, "Logit Value" if False.
        """
        if isinstance(logits, torch.Tensor):
            logits = logits.squeeze().detach().cpu().numpy()

        pos_logits = logits[position]

        if use_softmax:
            exp_logits = np.exp(pos_logits - np.max(pos_logits))
            pos_logits = exp_logits / np.sum(exp_logits)

        top_idx = np.argsort(pos_logits)[-k:][::-1]
        top_vals = pos_logits[top_idx]
        top_tokens = [self.tokenizer.decode([idx]) for idx in top_idx]

        plt.figure(figsize=(12, 6))
        plt.bar(range(k), top_vals)
        plt.xticks(range(k), top_tokens, rotation=45, ha='right', fontsize=10)
        plt.xlabel('Token', fontsize=14)
        ylabel = r'$P$(token)' if use_softmax else 'Logit Value'
        plt.ylabel(ylabel, fontsize=14)
        plt.title(title, fontsize=16)
        plt.tight_layout()
        if save_name is not None:
            plt.savefig(f"{save_name}.png", dpi=300, bbox_inches="tight")
        plt.show()

    def get_top_k_tokens(
            self,
            logits: torch.Tensor, 
            k: int = 10, 
            to_dataframe: bool = True,
            verbose: bool = False
        ) -> pd.DataFrame | list[dict]:
        """Get the top-k predicted tokens and their probabilities for each position in a sequence.

        Iterates through all positions in the sequence and extracts the k highest-scoring tokens
        at each position. Useful for detailed analysis of model predictions across the entire
        sequence.

        Args:
            logits (torch.Tensor): Logits tensor with shape (sequence_length, vocab_size)
                or similar. Will be automatically squeezed and moved to CPU if needed.
            k (int, optional): Number of top predictions to display per position.
                Defaults to 10.
            to_dataframe (bool, optional): If True, returns the top-k predicted tokens at each 
                position and their probabilities in a pandas DataFrame. Defaults to True.
            verbose (bool, optional): If True, prints the results to the console. Defaults
                to False. 

        Returns:
            pd.DataFrame | list[dict]: If to_dataframe=True, returns a DataFrame with columns
                'position', 'rank', 'token', and 'probability'. Otherwise returns a list of
                dictionaries with the same fields.

        Example Output:
            Top-10 predicted tokens per position

            Position 0:
                Token ('the'): 12.45
                Token ('a'): 11.32
                ...
        """
        if isinstance(logits, torch.Tensor):
            logits = logits.squeeze().detach().cpu().numpy()
        if verbose:
            print(f"Top-{k} predicted tokens per position")

        exp_logits = np.exp(logits - np.max(logits))
        logits = exp_logits / np.sum(exp_logits)

        results = {}
        for pos in range(logits.shape[0]):
            top_idx = np.argsort(logits[pos])[-k:][::-1]
            top_vals = logits[pos][top_idx]
            if verbose:
                print(f'\nPosition {pos}:')
            tokens = []
            for idx, val in zip(top_idx, top_vals):
                if self.tokenizer is not None:
                    token = self.tokenizer.decode([idx])
                    tokens.append(token)
                    if verbose:
                        print(f"\t'{token}': {val:.2f}")
                else:
                    if verbose:
                        print(f"\tToken {idx}: {val:.2f}")

            results[pos] = {
                'tokens': tokens,
                'values': top_vals.tolist()
            }

        out = []
        for position, data in results.items():
            for i, (token, prob) in enumerate(zip(data['tokens'], data['values'])):
                out.append({
                    'position': position,
                    'rank': i + 1,
                    'token': token,
                    'probability': prob
                })

        if to_dataframe:
            return pd.DataFrame(out)
        else:
            return out

    def get_most_active_features(
            self,
            sentences: list[str], 
            target: int | str | None = None,
            k: int | None = None,
            case_sensitive: bool = True,
            return_values: bool = False
        ) -> dict[str, FeatureResult]:
        """Extract SAE latent features for specific tokens across multiple sentences.

        This method processes a list of sentences through a transformer model and sparse
        autoencoder to extract active features at specified token positions. Useful for
        analyzing which SAE features activate for particular tokens or syntactic patterns.

        Note:
            Requires `hf_model`, `layer`, `sae_model`, and `sae_config` to be set during
            class initialization.

        Args:
            sentences (list[str]): List of input sentences to process. Each sentence will
                be tokenized and processed independently.
            target (int | str | None, optional): Specifies which token positions to analyze.
                - If int: Token position index (supports negative indexing, e.g., -1 for last token).
                - If str: Token string to match (e.g., "What"). Will find all occurrences.
                - If None: Extracts features for all token positions.
                Defaults to None.
            k (int | None, optional): If provided, returns only the top-k most active features
                by activation magnitude. If None, returns all non-zero features. Defaults to None.
            case_sensitive (bool, optional): Whether string matching for target tokens should
                be case-sensitive. Only applies when target is a string. Defaults to True.
            return_values (bool, optional): Whether the function will return features and activation
                values of each feature. Defaults to False.

        Returns:
            dict[str, FeatureResult]: Dictionary mapping descriptive keys to FeatureResult objects.
                Keys follow the format "sent_{idx}_pos_{pos}_tok_{token}" where:
                - idx: 1-indexed sentence number
                - pos: 0-indexed token position within the sentence
                - token: The decoded token string (stripped of whitespace)
                Values are FeatureResult dataclass instances with:
                - features: Tensor containing indices of active/top-k features
                - values: Tensor of activation values (or None if return_values=False)
        """
        features = {}
        for idx, sent in enumerate(sentences):
            acts = self.mlp_extractor.get_mlp_acts(sample=sent)
            input_ids = self.tokenizer.encode(sent, add_special_tokens=False)
            num_tokens = acts.shape[0]

            if target is None:
                positions = list(range(num_tokens))
            elif isinstance(target, int):
                pos = target if target >= 0 else num_tokens + target
                if 0 <= pos < num_tokens:
                    positions = [pos]
                else:
                    raise ValueError(
                        f"Position {target} out of range for sentence {idx+1} ({num_tokens} tokens)"
                    )
            elif isinstance(target, str):
                target_to_match = target if case_sensitive else target.lower()
                positions = []
                for i, tok_id in enumerate(input_ids):
                    decoded = self.tokenizer.decode([tok_id])
                    decoded_to_match = decoded if case_sensitive else decoded.lower()
                    if decoded_to_match.strip() == target_to_match.strip():
                        positions.append(i)
                if not positions:
                    print(f"Warning: '{target}' not found in sentence '{sent}'")
                    continue
            else:
                raise TypeError(f"target must be int, str, or None, got {type(target)}")

            for pos in positions:
                if pos >= num_tokens:
                    continue

                vals = None
                if return_values:
                    feats, vals = self.sae_extractor.get_alive_features(
                        activations=acts, 
                        token_position=pos, 
                        k=k,
                        return_values=True
                    )
                else:
                    feats = self.sae_extractor.get_alive_features(
                        activations=acts, 
                        token_position=pos, 
                        k=k,
                        return_values=False
                    )
                token_str = self.tokenizer.decode([input_ids[pos]])
                key = f"sent_{idx+1}_pos_{pos}_tok_{token_str.strip()}"
                features[key] = FeatureResult(
                    feats.cpu(), 
                    vals.cpu() if vals is not None else None
                )
        return features

`init(hf_model=None, sae_model=None, sae_config=None, layer=None, cache_dir='cache')` ¶

Initialize the AnalysisUtils class for analyzing sparse autoencoder results.

Parameters:

Name	Type	Description	Default
`hf_model`	`str`	Name or path of the HuggingFace model for extracting MLP activations (e.g., "gpt2", "meta-llama/Llama-2-7b"). Required for `get_most_active_features`. Defaults to None.	`None`
`sae_model`	`str`	Path to the trained sparse autoencoder model weights file. Required for `get_most_active_features`. Defaults to None.	`None`
`sae_config`	`str \| dict`	Path to the YAML configuration file or a dictionary containing SAE hyperparameters. Required for `get_most_active_features`. Defaults to None.	`None`
`layer`	`int`	Index of the transformer layer to extract activations from (0-indexed). Required for `get_most_active_features`. Defaults to None.	`None`
`cache_dir`	`str`	Directory to cache downloaded models. Defaults to 'cache'.	`'cache'`

Source code in deeplens\utils\analysis.py

def __init__(
        self,
        hf_model: str = None,
        sae_model: str = None,
        sae_config: str | dict = None,
        layer: int = None,
        cache_dir: str = 'cache'
    ):
    """Initialize the AnalysisUtils class for analyzing sparse autoencoder results.

    Args:
        hf_model (str, optional): Name or path of the HuggingFace model for extracting
            MLP activations (e.g., "gpt2", "meta-llama/Llama-2-7b"). Required for
            `get_most_active_features`. Defaults to None.
        sae_model (str, optional): Path to the trained sparse autoencoder model weights
            file. Required for `get_most_active_features`. Defaults to None.
        sae_config (str | dict, optional): Path to the YAML configuration file or a
            dictionary containing SAE hyperparameters. Required for 
            `get_most_active_features`. Defaults to None.
        layer (int, optional): Index of the transformer layer to extract activations
            from (0-indexed). Required for `get_most_active_features`. Defaults to None.
        cache_dir (str, optional): Directory to cache downloaded models.
            Defaults to 'cache'.
    """
    self.tokenizer = AutoTokenizer.from_pretrained(hf_model, cache_dir=cache_dir)
    self.hf_model = hf_model
    self.sae_model = sae_model
    self.sae_config = sae_config
    self.layer = layer

    self._mlp_extractor = None
    self._sae_extractor = None

`generate_logit_heatmap(logits, save_name=None, k=None)` ¶

Generate and display a heatmap visualization of logits across token positions.

Creates a color-coded heatmap showing the distribution of logit values across the vocabulary for each token position in a sequence. Uses the 'inferno' colormap for visualization.

Parameters:

Name	Type	Description	Default
`logits`	`Tensor`	Logits tensor with shape (sequence_length, vocab_size) or similar. Will be automatically squeezed and moved to CPU if needed.	required
`save_name`	`str`	Filename (without extension) to save the plot. If provided, saves the figure as a PNG file with 300 DPI. If None, only displays the plot. Defaults to None.	`None`
`k`	`int`	top k vocabulary indexes to plot. Defaults to None.	`None`

Returns:

Name	Type	Description
`None`	`None`	Displays the plot and optionally saves it to disk.

Note

The figure size is set to (20, 6) for optimal visualization of typical sequence lengths and vocabulary sizes.

Source code in deeplens\utils\analysis.py

def generate_logit_heatmap(
        self,
        logits: torch.Tensor, 
        save_name: str = None,
        k: int | None = None
    ) -> None:
    """Generate and display a heatmap visualization of logits across token positions.

    Creates a color-coded heatmap showing the distribution of logit values across the
    vocabulary for each token position in a sequence. Uses the 'inferno' colormap for
    visualization.

    Args:
        logits (torch.Tensor): Logits tensor with shape (sequence_length, vocab_size)
            or similar. Will be automatically squeezed and moved to CPU if needed.
        save_name (str, optional): Filename (without extension) to save the plot.
            If provided, saves the figure as a PNG file with 300 DPI. If None, only
            displays the plot. Defaults to None.
        k (int): top k vocabulary indexes to plot. Defaults to None.

    Returns:
        None: Displays the plot and optionally saves it to disk.

    Note:
        The figure size is set to (20, 6) for optimal visualization of typical
        sequence lengths and vocabulary sizes.
    """
    if isinstance(logits, torch.Tensor):
        logits = logits.squeeze().detach().cpu().numpy()

    if k is not None:
        if isinstance(logits, np.ndarray):
            logits = torch.from_numpy(logits)
        logits = torch.topk(logits, k=k, dim=-1).values.numpy()

    plt.figure(figsize=(20, 6))
    plt.imshow(logits, cmap="inferno", aspect="auto")
    plt.colorbar()
    plt.xlabel("Vocabulary Index", size=18)
    plt.ylabel("Token Position", size=18)
    plt.xticks(size=12)
    plt.yticks(size=12)
    if save_name is not None:
        plt.savefig(f"{save_name}.png", dpi=300, bbox_inches="tight")
    plt.show()

`get_most_active_features(sentences, target=None, k=None, case_sensitive=True, return_values=False)` ¶

Extract SAE latent features for specific tokens across multiple sentences.

This method processes a list of sentences through a transformer model and sparse autoencoder to extract active features at specified token positions. Useful for analyzing which SAE features activate for particular tokens or syntactic patterns.

Note

Requires hf_model, layer, sae_model, and sae_config to be set during class initialization.

Parameters:

Name	Type	Description	Default
`sentences`	`list[str]`	List of input sentences to process. Each sentence will be tokenized and processed independently.	required
`target`	`int \| str \| None`	Specifies which token positions to analyze. - If int: Token position index (supports negative indexing, e.g., -1 for last token). - If str: Token string to match (e.g., "What"). Will find all occurrences. - If None: Extracts features for all token positions. Defaults to None.	`None`
`k`	`int \| None`	If provided, returns only the top-k most active features by activation magnitude. If None, returns all non-zero features. Defaults to None.	`None`
`case_sensitive`	`bool`	Whether string matching for target tokens should be case-sensitive. Only applies when target is a string. Defaults to True.	`True`
`return_values`	`bool`	Whether the function will return features and activation values of each feature. Defaults to False.	`False`

Returns:

Type Description

dict[str, FeatureResult]

dict[str, FeatureResult]: Dictionary mapping descriptive keys to FeatureResult objects. Keys follow the format "sent_{idx}postok" where: - idx: 1-indexed sentence number - pos: 0-indexed token position within the sentence - token: The decoded token string (stripped of whitespace) Values are FeatureResult dataclass instances with: - features: Tensor containing indices of active/top-k features - values: Tensor of activation values (or None if return_values=False)

Source code in deeplens\utils\analysis.py

def get_most_active_features(
        self,
        sentences: list[str], 
        target: int | str | None = None,
        k: int | None = None,
        case_sensitive: bool = True,
        return_values: bool = False
    ) -> dict[str, FeatureResult]:
    """Extract SAE latent features for specific tokens across multiple sentences.

    This method processes a list of sentences through a transformer model and sparse
    autoencoder to extract active features at specified token positions. Useful for
    analyzing which SAE features activate for particular tokens or syntactic patterns.

    Note:
        Requires `hf_model`, `layer`, `sae_model`, and `sae_config` to be set during
        class initialization.

    Args:
        sentences (list[str]): List of input sentences to process. Each sentence will
            be tokenized and processed independently.
        target (int | str | None, optional): Specifies which token positions to analyze.
            - If int: Token position index (supports negative indexing, e.g., -1 for last token).
            - If str: Token string to match (e.g., "What"). Will find all occurrences.
            - If None: Extracts features for all token positions.
            Defaults to None.
        k (int | None, optional): If provided, returns only the top-k most active features
            by activation magnitude. If None, returns all non-zero features. Defaults to None.
        case_sensitive (bool, optional): Whether string matching for target tokens should
            be case-sensitive. Only applies when target is a string. Defaults to True.
        return_values (bool, optional): Whether the function will return features and activation
            values of each feature. Defaults to False.

    Returns:
        dict[str, FeatureResult]: Dictionary mapping descriptive keys to FeatureResult objects.
            Keys follow the format "sent_{idx}_pos_{pos}_tok_{token}" where:
            - idx: 1-indexed sentence number
            - pos: 0-indexed token position within the sentence
            - token: The decoded token string (stripped of whitespace)
            Values are FeatureResult dataclass instances with:
            - features: Tensor containing indices of active/top-k features
            - values: Tensor of activation values (or None if return_values=False)
    """
    features = {}
    for idx, sent in enumerate(sentences):
        acts = self.mlp_extractor.get_mlp_acts(sample=sent)
        input_ids = self.tokenizer.encode(sent, add_special_tokens=False)
        num_tokens = acts.shape[0]

        if target is None:
            positions = list(range(num_tokens))
        elif isinstance(target, int):
            pos = target if target >= 0 else num_tokens + target
            if 0 <= pos < num_tokens:
                positions = [pos]
            else:
                raise ValueError(
                    f"Position {target} out of range for sentence {idx+1} ({num_tokens} tokens)"
                )
        elif isinstance(target, str):
            target_to_match = target if case_sensitive else target.lower()
            positions = []
            for i, tok_id in enumerate(input_ids):
                decoded = self.tokenizer.decode([tok_id])
                decoded_to_match = decoded if case_sensitive else decoded.lower()
                if decoded_to_match.strip() == target_to_match.strip():
                    positions.append(i)
            if not positions:
                print(f"Warning: '{target}' not found in sentence '{sent}'")
                continue
        else:
            raise TypeError(f"target must be int, str, or None, got {type(target)}")

        for pos in positions:
            if pos >= num_tokens:
                continue

            vals = None
            if return_values:
                feats, vals = self.sae_extractor.get_alive_features(
                    activations=acts, 
                    token_position=pos, 
                    k=k,
                    return_values=True
                )
            else:
                feats = self.sae_extractor.get_alive_features(
                    activations=acts, 
                    token_position=pos, 
                    k=k,
                    return_values=False
                )
            token_str = self.tokenizer.decode([input_ids[pos]])
            key = f"sent_{idx+1}_pos_{pos}_tok_{token_str.strip()}"
            features[key] = FeatureResult(
                feats.cpu(), 
                vals.cpu() if vals is not None else None
            )
    return features

`get_top_k_tokens(logits, k=10, to_dataframe=True, verbose=False)` ¶

Get the top-k predicted tokens and their probabilities for each position in a sequence.

Iterates through all positions in the sequence and extracts the k highest-scoring tokens at each position. Useful for detailed analysis of model predictions across the entire sequence.

Parameters:

Name	Type	Description	Default
`logits`	`Tensor`	Logits tensor with shape (sequence_length, vocab_size) or similar. Will be automatically squeezed and moved to CPU if needed.	required
`k`	`int`	Number of top predictions to display per position. Defaults to 10.	`10`
`to_dataframe`	`bool`	If True, returns the top-k predicted tokens at each position and their probabilities in a pandas DataFrame. Defaults to True.	`True`
`verbose`	`bool`	If True, prints the results to the console. Defaults to False.	`False`

Returns:

Type	Description
`DataFrame \| list[dict]`	pd.DataFrame \| list[dict]: If to_dataframe=True, returns a DataFrame with columns 'position', 'rank', 'token', and 'probability'. Otherwise returns a list of dictionaries with the same fields.

Example Output

Top-10 predicted tokens per position

Position 0: Token ('the'): 12.45 Token ('a'): 11.32 ...

Source code in deeplens\utils\analysis.py

def get_top_k_tokens(
        self,
        logits: torch.Tensor, 
        k: int = 10, 
        to_dataframe: bool = True,
        verbose: bool = False
    ) -> pd.DataFrame | list[dict]:
    """Get the top-k predicted tokens and their probabilities for each position in a sequence.

    Iterates through all positions in the sequence and extracts the k highest-scoring tokens
    at each position. Useful for detailed analysis of model predictions across the entire
    sequence.

    Args:
        logits (torch.Tensor): Logits tensor with shape (sequence_length, vocab_size)
            or similar. Will be automatically squeezed and moved to CPU if needed.
        k (int, optional): Number of top predictions to display per position.
            Defaults to 10.
        to_dataframe (bool, optional): If True, returns the top-k predicted tokens at each 
            position and their probabilities in a pandas DataFrame. Defaults to True.
        verbose (bool, optional): If True, prints the results to the console. Defaults
            to False. 

    Returns:
        pd.DataFrame | list[dict]: If to_dataframe=True, returns a DataFrame with columns
            'position', 'rank', 'token', and 'probability'. Otherwise returns a list of
            dictionaries with the same fields.

    Example Output:
        Top-10 predicted tokens per position

        Position 0:
            Token ('the'): 12.45
            Token ('a'): 11.32
            ...
    """
    if isinstance(logits, torch.Tensor):
        logits = logits.squeeze().detach().cpu().numpy()
    if verbose:
        print(f"Top-{k} predicted tokens per position")

    exp_logits = np.exp(logits - np.max(logits))
    logits = exp_logits / np.sum(exp_logits)

    results = {}
    for pos in range(logits.shape[0]):
        top_idx = np.argsort(logits[pos])[-k:][::-1]
        top_vals = logits[pos][top_idx]
        if verbose:
            print(f'\nPosition {pos}:')
        tokens = []
        for idx, val in zip(top_idx, top_vals):
            if self.tokenizer is not None:
                token = self.tokenizer.decode([idx])
                tokens.append(token)
                if verbose:
                    print(f"\t'{token}': {val:.2f}")
            else:
                if verbose:
                    print(f"\tToken {idx}: {val:.2f}")

        results[pos] = {
            'tokens': tokens,
            'values': top_vals.tolist()
        }

    out = []
    for position, data in results.items():
        for i, (token, prob) in enumerate(zip(data['tokens'], data['values'])):
            out.append({
                'position': position,
                'rank': i + 1,
                'token': token,
                'probability': prob
            })

    if to_dataframe:
        return pd.DataFrame(out)
    else:
        return out

`plot_topk_distribution(logits, k=20, position=0, save_name=None, use_softmax=True, title=None)` ¶

Plot a bar chart showing the top-k most probable tokens at a specific sequence position.

Creates a horizontal bar chart displaying the highest-probability token predictions at a given position in the sequence. Useful for analyzing model predictions and understanding which tokens are most likely at specific positions.

Parameters:

Name	Type	Description	Default
`logits`	`Tensor`	Logits tensor with shape (sequence_length, vocab_size) or similar. Will be automatically squeezed and moved to CPU if needed.	required
`k`	`int`	Number of top predictions to display in the bar chart. Defaults to 20.	`20`
`position`	`int`	Token position in the sequence to analyze. Use 0-based indexing. Defaults to 0 (first token).	`0`
`save_name`	`str`	Filename (without extension) to save the plot. If provided, saves the figure as a PNG file with 300 DPI. If None, only displays the plot. Defaults to None.	`None`
`use_softmax`	`bool`	If True, applies softmax to convert logits to probabilities before plotting. If False, plots raw logit values. Defaults to True.	`True`
`title`	`str`	Custom title for the plot. If None, no title is displayed. Defaults to None.	`None`

Returns:

Name	Type	Description
`None`	`None`	Displays the plot and optionally saves it to disk.

Note

Token labels are rotated 45 degrees for readability. The y-axis label changes based on use_softmax: shows probability notation if True, "Logit Value" if False.

Source code in deeplens\utils\analysis.py

def plot_topk_distribution(
        self,
        logits: torch.Tensor, 
        k: int = 20, 
        position: int = 0, 
        save_name: str = None,
        use_softmax: bool = True,
        title: str = None
    ) -> None:
    """Plot a bar chart showing the top-k most probable tokens at a specific sequence position.

    Creates a horizontal bar chart displaying the highest-probability token predictions
    at a given position in the sequence. Useful for analyzing model predictions and
    understanding which tokens are most likely at specific positions.

    Args:
        logits (torch.Tensor): Logits tensor with shape (sequence_length, vocab_size)
            or similar. Will be automatically squeezed and moved to CPU if needed.
        k (int, optional): Number of top predictions to display in the bar chart.
            Defaults to 20.
        position (int, optional): Token position in the sequence to analyze. Use 0-based
            indexing. Defaults to 0 (first token).
        save_name (str, optional): Filename (without extension) to save the plot.
            If provided, saves the figure as a PNG file with 300 DPI. If None, only
            displays the plot. Defaults to None.
        use_softmax (bool, optional): If True, applies softmax to convert logits to
            probabilities before plotting. If False, plots raw logit values.
            Defaults to True.
        title (str, optional): Custom title for the plot. If None, no title is displayed.
            Defaults to None.

    Returns:
        None: Displays the plot and optionally saves it to disk.

    Note:
        Token labels are rotated 45 degrees for readability. The y-axis label changes
        based on use_softmax: shows probability notation if True, "Logit Value" if False.
    """
    if isinstance(logits, torch.Tensor):
        logits = logits.squeeze().detach().cpu().numpy()

    pos_logits = logits[position]

    if use_softmax:
        exp_logits = np.exp(pos_logits - np.max(pos_logits))
        pos_logits = exp_logits / np.sum(exp_logits)

    top_idx = np.argsort(pos_logits)[-k:][::-1]
    top_vals = pos_logits[top_idx]
    top_tokens = [self.tokenizer.decode([idx]) for idx in top_idx]

    plt.figure(figsize=(12, 6))
    plt.bar(range(k), top_vals)
    plt.xticks(range(k), top_tokens, rotation=45, ha='right', fontsize=10)
    plt.xlabel('Token', fontsize=14)
    ylabel = r'$P$(token)' if use_softmax else 'Logit Value'
    plt.ylabel(ylabel, fontsize=14)
    plt.title(title, fontsize=16)
    plt.tight_layout()
    if save_name is not None:
        plt.savefig(f"{save_name}.png", dpi=300, bbox_inches="tight")
    plt.show()

`deeplens.utils.dataset` ¶

`ActivationsDataset` ¶

Bases: Dataset

Lightweight PyTorch Dataset wrapper for pre-computed activation tensors.

Provides a simple Dataset interface for tensors of neural network activations, enabling use with PyTorch DataLoader for batching and iteration.

Source code in deeplens\utils\dataset.py

class ActivationsDataset(Dataset):
    """Lightweight PyTorch Dataset wrapper for pre-computed activation tensors.

    Provides a simple Dataset interface for tensors of neural network activations,
    enabling use with PyTorch DataLoader for batching and iteration.
    """
    def __init__(self, activations: torch.Tensor):
        """Initialize the dataset with activation tensors.

        Args:
            activations (torch.Tensor): Tensor containing pre-computed activations
                with shape (num_samples, feature_dim).
        """
        super().__init__()
        self.activations = activations

    def __len__(self) -> int:
        """Get the number of activation samples.

        Returns:
            int: Number of samples in the dataset.
        """
        return len(self.activations)

    def __getitem__(self, idx) -> torch.Tensor:
        """Retrieve activation tensor at the specified index.

        Args:
            idx (int): Index of the sample to retrieve.

        Returns:
            torch.Tensor: Activation tensor at the given index.
        """
        return self.activations[idx]

`getitem(idx)` ¶

Retrieve activation tensor at the specified index.

Parameters:

Name	Type	Description	Default
`idx`	`int`	Index of the sample to retrieve.	required

Returns:

Type	Description
`Tensor`	torch.Tensor: Activation tensor at the given index.

Source code in deeplens\utils\dataset.py

def __getitem__(self, idx) -> torch.Tensor:
    """Retrieve activation tensor at the specified index.

    Args:
        idx (int): Index of the sample to retrieve.

    Returns:
        torch.Tensor: Activation tensor at the given index.
    """
    return self.activations[idx]

`init(activations)` ¶

Initialize the dataset with activation tensors.

Parameters:

Name	Type	Description	Default
`activations`	`Tensor`	Tensor containing pre-computed activations with shape (num_samples, feature_dim).	required

Source code in deeplens\utils\dataset.py

def __init__(self, activations: torch.Tensor):
    """Initialize the dataset with activation tensors.

    Args:
        activations (torch.Tensor): Tensor containing pre-computed activations
            with shape (num_samples, feature_dim).
    """
    super().__init__()
    self.activations = activations

`len()` ¶

Get the number of activation samples.

Returns:

Name	Type	Description
`int`	`int`	Number of samples in the dataset.

Source code in deeplens\utils\dataset.py

def __len__(self) -> int:
    """Get the number of activation samples.

    Returns:
        int: Number of samples in the dataset.
    """
    return len(self.activations)

`ActivationsDatasetBuilder` ¶

Builder class for creating DataLoaders from saved activation tensors.

Loads pre-computed activations from disk, applies optional normalization, and creates train/validation DataLoaders for training sparse autoencoders or other downstream models.

Source code in deeplens\utils\dataset.py

class ActivationsDatasetBuilder():
    """Builder class for creating DataLoaders from saved activation tensors.

    Loads pre-computed activations from disk, applies optional normalization, and creates
    train/validation DataLoaders for training sparse autoencoders or other downstream models.
    """
    def __init__(
            self, 
            activations: torch.Tensor = None, 
            splits: list = [0.8, 0.2],
            batch_size: int = 16,
            norm: bool = True
        ):
        """Initialize the builder and load activations from disk.

        Args:
            activations (torch.Tensor | str, optional): Path to a .pt file containing
                saved activation tensors, or a tensor directly. Defaults to None.
            splits (list, optional): List of two floats representing train and validation
                split proportions. Must sum to 1.0. Defaults to [0.8, 0.2].
            batch_size (int, optional): Number of samples per batch for DataLoaders.
                Defaults to 16.
            norm (bool, optional): Whether to apply z-score normalization (standardization)
                to the activations. Defaults to True.
        """
        self.activations = torch.load(activations, weights_only=True)
        self.splits = splits
        self.batch_size = batch_size
        self.norm = norm
        self.normalize()

    def set_tensor_dataset(self) -> Dataset:
        """Create a PyTorch Dataset from the loaded activations.

        Returns:
            Dataset: ActivationsDataset instance wrapping the activation tensors.
        """
        return ActivationsDataset(self.activations)

    def get_dataloaders(self, ddp: bool = False) -> tuple[DataLoader, DataLoader]:
        """Create train and validation DataLoaders from the activations.

        Splits the dataset according to the specified proportions and creates two DataLoaders
        with appropriate settings for training and evaluation.

        Args:
            ddp (bool, optional): Turn to True if Distributed Data Parallel (DDP) training is
                intended. Defaults to False.

        Returns:
            tuple: A tuple containing (train_loader, eval_loader).
                Training loader has shuffle=True for randomized batching, evaluation loader
                has shuffle=False for consistent evaluation.
        """
        data = self.set_tensor_dataset()
        train, eval = random_split(data, lengths=self.splits)
        train_loader = DataLoader(
            train, 
            batch_size=self.batch_size, 
            shuffle=not ddp, 
            pin_memory=True,
            sampler=DistributedSampler(train, shuffle=True) if ddp else None
        )
        eval_loader = DataLoader(
            eval, 
            batch_size=self.batch_size, 
            shuffle=False, 
            pin_memory=True, 
            sampler=DistributedSampler(eval, shuffle=False) if ddp else None
        )
        return train_loader, eval_loader

    @torch.no_grad()
    def normalize(self) -> None:
        """Apply z-score normalization to the activations in-place.

        Standardizes the activations by subtracting the mean and dividing by the standard
        deviation (computed along the batch dimension). Adds small epsilon (1e-8) to prevent
        division by zero. Only applies if norm=True was set during initialization.

        Returns:
            None: Modifies self.activations in-place.
        """
        if self.norm:
            mean = self.activations.mean(dim=0, keepdim=True)
            std = self.activations.std(dim=0, keepdim=True)
            self.activations = (self.activations - mean) / (std + 1e-8)

`init(activations=None, splits=[0.8, 0.2], batch_size=16, norm=True)` ¶

Initialize the builder and load activations from disk.

Parameters:

Name	Type	Description	Default
`activations`	`Tensor \| str`	Path to a .pt file containing saved activation tensors, or a tensor directly. Defaults to None.	`None`
`splits`	`list`	List of two floats representing train and validation split proportions. Must sum to 1.0. Defaults to [0.8, 0.2].	`[0.8, 0.2]`
`batch_size`	`int`	Number of samples per batch for DataLoaders. Defaults to 16.	`16`
`norm`	`bool`	Whether to apply z-score normalization (standardization) to the activations. Defaults to True.	`True`

Source code in deeplens\utils\dataset.py

def __init__(
        self, 
        activations: torch.Tensor = None, 
        splits: list = [0.8, 0.2],
        batch_size: int = 16,
        norm: bool = True
    ):
    """Initialize the builder and load activations from disk.

    Args:
        activations (torch.Tensor | str, optional): Path to a .pt file containing
            saved activation tensors, or a tensor directly. Defaults to None.
        splits (list, optional): List of two floats representing train and validation
            split proportions. Must sum to 1.0. Defaults to [0.8, 0.2].
        batch_size (int, optional): Number of samples per batch for DataLoaders.
            Defaults to 16.
        norm (bool, optional): Whether to apply z-score normalization (standardization)
            to the activations. Defaults to True.
    """
    self.activations = torch.load(activations, weights_only=True)
    self.splits = splits
    self.batch_size = batch_size
    self.norm = norm
    self.normalize()

`get_dataloaders(ddp=False)` ¶

Create train and validation DataLoaders from the activations.

Splits the dataset according to the specified proportions and creates two DataLoaders with appropriate settings for training and evaluation.

Parameters:

Name	Type	Description	Default
`ddp`	`bool`	Turn to True if Distributed Data Parallel (DDP) training is intended. Defaults to False.	`False`

Returns:

Name	Type	Description
`tuple`	`tuple[DataLoader, DataLoader]`	A tuple containing (train_loader, eval_loader). Training loader has shuffle=True for randomized batching, evaluation loader has shuffle=False for consistent evaluation.

Source code in deeplens\utils\dataset.py

def get_dataloaders(self, ddp: bool = False) -> tuple[DataLoader, DataLoader]:
    """Create train and validation DataLoaders from the activations.

    Splits the dataset according to the specified proportions and creates two DataLoaders
    with appropriate settings for training and evaluation.

    Args:
        ddp (bool, optional): Turn to True if Distributed Data Parallel (DDP) training is
            intended. Defaults to False.

    Returns:
        tuple: A tuple containing (train_loader, eval_loader).
            Training loader has shuffle=True for randomized batching, evaluation loader
            has shuffle=False for consistent evaluation.
    """
    data = self.set_tensor_dataset()
    train, eval = random_split(data, lengths=self.splits)
    train_loader = DataLoader(
        train, 
        batch_size=self.batch_size, 
        shuffle=not ddp, 
        pin_memory=True,
        sampler=DistributedSampler(train, shuffle=True) if ddp else None
    )
    eval_loader = DataLoader(
        eval, 
        batch_size=self.batch_size, 
        shuffle=False, 
        pin_memory=True, 
        sampler=DistributedSampler(eval, shuffle=False) if ddp else None
    )
    return train_loader, eval_loader

`normalize()` ¶

Apply z-score normalization to the activations in-place.

Standardizes the activations by subtracting the mean and dividing by the standard deviation (computed along the batch dimension). Adds small epsilon (1e-8) to prevent division by zero. Only applies if norm=True was set during initialization.

Returns:

Name	Type	Description
`None`	`None`	Modifies self.activations in-place.

Source code in deeplens\utils\dataset.py

@torch.no_grad()
def normalize(self) -> None:
    """Apply z-score normalization to the activations in-place.

    Standardizes the activations by subtracting the mean and dividing by the standard
    deviation (computed along the batch dimension). Adds small epsilon (1e-8) to prevent
    division by zero. Only applies if norm=True was set during initialization.

    Returns:
        None: Modifies self.activations in-place.
    """
    if self.norm:
        mean = self.activations.mean(dim=0, keepdim=True)
        std = self.activations.std(dim=0, keepdim=True)
        self.activations = (self.activations - mean) / (std + 1e-8)

`set_tensor_dataset()` ¶

Create a PyTorch Dataset from the loaded activations.

Returns:

Name	Type	Description
`Dataset`	`Dataset`	ActivationsDataset instance wrapping the activation tensors.

Source code in deeplens\utils\dataset.py

def set_tensor_dataset(self) -> Dataset:
    """Create a PyTorch Dataset from the loaded activations.

    Returns:
        Dataset: ActivationsDataset instance wrapping the activation tensors.
    """
    return ActivationsDataset(self.activations)

`AudioDatasetBuilder` ¶

Bases: Dataset

PyTorch Dataset for loading and preprocessing audio files with optional mel spectrogram transformation.

This dataset handles audio files in various formats (WAV, MP3, FLAC), performs preprocessing operations like resampling, mono conversion, and padding, and optionally converts waveforms to mel spectrograms. Supports both labeled and unlabeled datasets.

Source code in deeplens\utils\dataset.py

class AudioDatasetBuilder(Dataset):
    """PyTorch Dataset for loading and preprocessing audio files with optional mel spectrogram transformation.

    This dataset handles audio files in various formats (WAV, MP3, FLAC), performs preprocessing
    operations like resampling, mono conversion, and padding, and optionally converts waveforms
    to mel spectrograms. Supports both labeled and unlabeled datasets.
    """
    def __init__(
            self, 
            audio_dir: str = None, 
            annotations_file: str = None, 
            target_sample_rate: int = 22050, 
            num_samples: int = 22050, 
            device: str = "auto",
            transformation_args: dict = {"n_fft": 1024, "hop_length": 512, "n_mels": 64}
        ) -> None:
        """Initialize the AudioDatasetBuilder with audio preprocessing parameters.

        Args:
            audio_dir (str, optional): Path to directory containing audio files. The directory
                should contain files with extensions .wav, .mp3, or .flac. Defaults to None.
            annotations_file (str, optional): Path to CSV file containing annotations/labels.
                If None, dataset returns only audio without labels. Defaults to None.
            target_sample_rate (int, optional): Target sampling rate in Hz for resampling.
                All audio will be resampled to this rate. Defaults to 22050.
            num_samples (int, optional): Target number of samples per audio clip. Audio will
                be truncated or zero-padded to this length. Defaults to 22050.
            device (str, optional): Device for tensor operations. Can be "auto" for automatic
                selection, "cuda", "mps", or "cpu". Defaults to "auto".
            transformation_args (dict, optional): Dictionary containing parameters for mel
                spectrogram transformation. Must include keys: "n_fft", "hop_length", "n_mels".
                If None, raw waveforms are returned without transformation. Defaults to
                {"n_fft": 1024, "hop_length": 512, "n_mels": 64}.
        """
        super().__init__()
        self.audio_dir = audio_dir 
        self.file_list = [
            f for f in os.listdir(self.audio_dir) 
            if f.endswith(('.wav', '.mp3', '.flac'))
        ]
        if annotations_file is not None:
            self.annotations = pd.read_csv(annotations_file)
        else:
            self.annotations = None

        self.target_sample_rate = target_sample_rate
        self.num_samples = num_samples

        if device == "auto":
            self.device = torch.device(
                "cuda" if torch.cuda.is_available() 
                else "mps" if torch.backends.mps.is_available()
                else "cpu"
            )     
        else:
            self.device = torch.device(device)

        self.transformation_args = transformation_args
        if transformation_args is not None:
            assert {"n_fft", "hop_length", "n_mels"}.issubset(transformation_args.keys()), \
            "Missing arguments. Please provide n_fft, hop_length, and n_mels."
            self.mel_spectrogram = torchaudio.transforms.MelSpectrogram(
                sample_rate=self.target_sample_rate,
                n_fft = transformation_args["n_fft"],
                hop_length=transformation_args["hop_length"],
                n_mels=transformation_args["n_mels"]
            ).to(self.device)
        else:
            self.mel_spectrogram = None

    def __len__(self) -> int:
        """Get the total number of samples in the dataset.

        Returns:
            int: Number of audio samples. Returns length of annotations if provided,
                otherwise returns number of audio files in the directory.
        """
        if self.annotations is not None:
            return len(self.annotations)
        else:
            return len(self.file_list)

    def __getitem__(self, index) -> torch.Tensor | tuple[torch.Tensor, int]:
        """Retrieve and preprocess an audio sample at the specified index.

        Loads the audio file, applies preprocessing (resampling, mono conversion, padding/truncation),
        optionally transforms to mel spectrogram, and returns with label if annotations are available.

        Args:
            index (int): Index of the sample to retrieve.

        Returns:
            torch.Tensor | tuple[torch.Tensor, int]: If annotations are provided, returns a tuple
                of (processed_audio, label). Otherwise, returns only the processed audio tensor.
                Audio shape depends on transformation: (1, num_samples) for waveform or
                (1, n_mels, time_steps) for mel spectrogram.
        """
        audio_sample = self._audio_sample_path(index)    
        signal, sr = torchaudio.load(audio_sample)
        signal = signal.to(self.device)                   
        signal = self._resample_if_necessary(signal, sr)  
        signal = self._mix_down_if_necessary(signal)  
        signal = self._truncate_if_necessary(signal)    
        signal = self._pad_if_necessary(signal)
        if self.transformation_args is not None:
            signal = self._apply_transformation(signal)
        if self.annotations is not None:
            label = self._audio_sample_label(index)        
            return signal, label
        else:
            return signal

    def _audio_sample_path(self, index) -> str:
        """Construct the file path for an audio sample at the given index.

        Args:
            index (int): Index of the audio sample.

        Returns:
            str: Full path to the audio file. If annotations are provided, constructs path
                using fold structure; otherwise, uses direct file listing.
        """
        if self.annotations is not None:
            fold = f"fold{self.annotations.iloc[index, 5]}"
            path = os.path.join(self.audio_dir, fold, self.annotations.iloc[index, 0])
        else:
            path = os.path.join(self.audio_dir, self.file_list[index])
        return path

    def _audio_sample_label(self, index) -> int:
        """Retrieve the label for an audio sample at the given index.

        Args:
            index (int): Index of the audio sample.

        Returns:
            int: Label value from the annotations file (column 6).
        """
        return self.annotations.iloc[index, 6]

    @torch.no_grad()
    def _resample_if_necessary(self, signal, sr) -> torch.Tensor:
        """Resample audio signal to target sample rate if necessary.

        Args:
            signal (torch.Tensor): Input audio waveform.
            sr (int): Current sample rate of the audio signal.

        Returns:
            torch.Tensor: Resampled audio at target_sample_rate, or original signal
                if already at the target rate.
        """
        if sr != self.target_sample_rate:
            resampler = torchaudio.transforms.Resample(sr, self.target_sample_rate)
            resampler.to(self.device)
            signal = resampler(signal)
        return signal

    @torch.no_grad()
    def _mix_down_if_necessary(self, signal) -> torch.Tensor:
        """Convert stereo audio to mono by averaging channels if necessary.

        Args:
            signal (torch.Tensor): Input audio with shape (channels, samples).

        Returns:
            torch.Tensor: Mono audio with shape (1, samples). If input is already mono,
                returns unchanged.
        """
        if signal.shape[0] > 1:
            signal = torch.mean(signal, dim=0, keepdim=True)
        return signal

    def _truncate_if_necessary(self, signal) -> torch.Tensor:
        """Truncate audio signal to target length if it exceeds num_samples.

        Args:
            signal (torch.Tensor): Input audio waveform.

        Returns:
            torch.Tensor: Truncated audio limited to num_samples length, or original
                signal if already shorter.
        """
        if signal.shape[1] > self.num_samples:
            signal = signal[:, :self.num_samples]
        return signal

    @torch.no_grad()
    def _pad_if_necessary(self, signal) -> torch.Tensor:
        """Zero-pad audio signal to target length if it's shorter than num_samples.

        Args:
            signal (torch.Tensor): Input audio waveform.

        Returns:
            torch.Tensor: Zero-padded audio extended to num_samples length, or original
                signal if already long enough.
        """
        length_signal = signal.shape[1]
        if length_signal < self.num_samples:
            n_padding = self.num_samples - length_signal
            r_pad_dim = (0, n_padding)
            signal = torch.nn.functional.pad(signal, r_pad_dim)
        return signal

    def _apply_transformation(self, signal) -> torch.Tensor:
        """Transform audio waveform to mel spectrogram representation.

        Args:
            signal (torch.Tensor): Input audio waveform with shape (1, num_samples).

        Returns:
            torch.Tensor: Mel spectrogram with shape (1, n_mels, time_steps), where
                time_steps depends on n_fft and hop_length parameters.
        """
        spectrogram = self.mel_spectrogram(signal)
        return spectrogram

`getitem(index)` ¶

Retrieve and preprocess an audio sample at the specified index.

Loads the audio file, applies preprocessing (resampling, mono conversion, padding/truncation), optionally transforms to mel spectrogram, and returns with label if annotations are available.

Parameters:

Name	Type	Description	Default
`index`	`int`	Index of the sample to retrieve.	required

Returns:

Type	Description
`Tensor \| tuple[Tensor, int]`	torch.Tensor \| tuple[torch.Tensor, int]: If annotations are provided, returns a tuple of (processed_audio, label). Otherwise, returns only the processed audio tensor. Audio shape depends on transformation: (1, num_samples) for waveform or (1, n_mels, time_steps) for mel spectrogram.

Source code in deeplens\utils\dataset.py

def __getitem__(self, index) -> torch.Tensor | tuple[torch.Tensor, int]:
    """Retrieve and preprocess an audio sample at the specified index.

    Loads the audio file, applies preprocessing (resampling, mono conversion, padding/truncation),
    optionally transforms to mel spectrogram, and returns with label if annotations are available.

    Args:
        index (int): Index of the sample to retrieve.

    Returns:
        torch.Tensor | tuple[torch.Tensor, int]: If annotations are provided, returns a tuple
            of (processed_audio, label). Otherwise, returns only the processed audio tensor.
            Audio shape depends on transformation: (1, num_samples) for waveform or
            (1, n_mels, time_steps) for mel spectrogram.
    """
    audio_sample = self._audio_sample_path(index)    
    signal, sr = torchaudio.load(audio_sample)
    signal = signal.to(self.device)                   
    signal = self._resample_if_necessary(signal, sr)  
    signal = self._mix_down_if_necessary(signal)  
    signal = self._truncate_if_necessary(signal)    
    signal = self._pad_if_necessary(signal)
    if self.transformation_args is not None:
        signal = self._apply_transformation(signal)
    if self.annotations is not None:
        label = self._audio_sample_label(index)        
        return signal, label
    else:
        return signal

`init(audio_dir=None, annotations_file=None, target_sample_rate=22050, num_samples=22050, device='auto', transformation_args={'n_fft': 1024, 'hop_length': 512, 'n_mels': 64})` ¶

Initialize the AudioDatasetBuilder with audio preprocessing parameters.

Parameters:

Name	Type	Description	Default
`audio_dir`	`str`	Path to directory containing audio files. The directory should contain files with extensions .wav, .mp3, or .flac. Defaults to None.	`None`
`annotations_file`	`str`	Path to CSV file containing annotations/labels. If None, dataset returns only audio without labels. Defaults to None.	`None`
`target_sample_rate`	`int`	Target sampling rate in Hz for resampling. All audio will be resampled to this rate. Defaults to 22050.	`22050`
`num_samples`	`int`	Target number of samples per audio clip. Audio will be truncated or zero-padded to this length. Defaults to 22050.	`22050`
`device`	`str`	Device for tensor operations. Can be "auto" for automatic selection, "cuda", "mps", or "cpu". Defaults to "auto".	`'auto'`
`transformation_args`	`dict`	Dictionary containing parameters for mel spectrogram transformation. Must include keys: "n_fft", "hop_length", "n_mels". If None, raw waveforms are returned without transformation. Defaults to {"n_fft": 1024, "hop_length": 512, "n_mels": 64}.	`{'n_fft': 1024, 'hop_length': 512, 'n_mels': 64}`

Source code in deeplens\utils\dataset.py

def __init__(
        self, 
        audio_dir: str = None, 
        annotations_file: str = None, 
        target_sample_rate: int = 22050, 
        num_samples: int = 22050, 
        device: str = "auto",
        transformation_args: dict = {"n_fft": 1024, "hop_length": 512, "n_mels": 64}
    ) -> None:
    """Initialize the AudioDatasetBuilder with audio preprocessing parameters.

    Args:
        audio_dir (str, optional): Path to directory containing audio files. The directory
            should contain files with extensions .wav, .mp3, or .flac. Defaults to None.
        annotations_file (str, optional): Path to CSV file containing annotations/labels.
            If None, dataset returns only audio without labels. Defaults to None.
        target_sample_rate (int, optional): Target sampling rate in Hz for resampling.
            All audio will be resampled to this rate. Defaults to 22050.
        num_samples (int, optional): Target number of samples per audio clip. Audio will
            be truncated or zero-padded to this length. Defaults to 22050.
        device (str, optional): Device for tensor operations. Can be "auto" for automatic
            selection, "cuda", "mps", or "cpu". Defaults to "auto".
        transformation_args (dict, optional): Dictionary containing parameters for mel
            spectrogram transformation. Must include keys: "n_fft", "hop_length", "n_mels".
            If None, raw waveforms are returned without transformation. Defaults to
            {"n_fft": 1024, "hop_length": 512, "n_mels": 64}.
    """
    super().__init__()
    self.audio_dir = audio_dir 
    self.file_list = [
        f for f in os.listdir(self.audio_dir) 
        if f.endswith(('.wav', '.mp3', '.flac'))
    ]
    if annotations_file is not None:
        self.annotations = pd.read_csv(annotations_file)
    else:
        self.annotations = None

    self.target_sample_rate = target_sample_rate
    self.num_samples = num_samples

    if device == "auto":
        self.device = torch.device(
            "cuda" if torch.cuda.is_available() 
            else "mps" if torch.backends.mps.is_available()
            else "cpu"
        )     
    else:
        self.device = torch.device(device)

    self.transformation_args = transformation_args
    if transformation_args is not None:
        assert {"n_fft", "hop_length", "n_mels"}.issubset(transformation_args.keys()), \
        "Missing arguments. Please provide n_fft, hop_length, and n_mels."
        self.mel_spectrogram = torchaudio.transforms.MelSpectrogram(
            sample_rate=self.target_sample_rate,
            n_fft = transformation_args["n_fft"],
            hop_length=transformation_args["hop_length"],
            n_mels=transformation_args["n_mels"]
        ).to(self.device)
    else:
        self.mel_spectrogram = None

`len()` ¶

Get the total number of samples in the dataset.

Returns:

Name	Type	Description
`int`	`int`	Number of audio samples. Returns length of annotations if provided, otherwise returns number of audio files in the directory.

Source code in deeplens\utils\dataset.py

def __len__(self) -> int:
    """Get the total number of samples in the dataset.

    Returns:
        int: Number of audio samples. Returns length of annotations if provided,
            otherwise returns number of audio files in the directory.
    """
    if self.annotations is not None:
        return len(self.annotations)
    else:
        return len(self.file_list)

`GetDataLoaders` ¶

Utility class for creating train and test DataLoaders from a PyTorch Dataset.

Handles dataset splitting and DataLoader creation with consistent parameters.

Source code in deeplens\utils\dataset.py

class GetDataLoaders():
    """Utility class for creating train and test DataLoaders from a PyTorch Dataset.

    Handles dataset splitting and DataLoader creation with consistent parameters.
    """
    def __init__(
            self, 
            dataset: Dataset = None,
            splits: list = [0.8, 0.2],
            batch_size: int = 16
        ) -> None:
        """Initialize the DataLoader factory.

        Args:
            dataset (Dataset, optional): PyTorch Dataset to split and load. Defaults to None.
            splits (list, optional): List of two floats representing train and test split
                proportions. Must sum to 1.0. Defaults to [0.8, 0.2].
            batch_size (int, optional): Number of samples per batch. Defaults to 16.
        """
        self.dataset = dataset
        self.splits = splits
        self.batch_size = batch_size

    def _prepare_loader(self) -> tuple[DataLoader, DataLoader]:
        """Create train and test DataLoaders with the specified configuration.

        Splits the dataset according to the split proportions and creates two DataLoaders
        with appropriate settings for training and testing.

        Returns:
            tuple[DataLoader, DataLoader]: A tuple containing (train_loader, test_loader).
                Training loader has shuffle=True, test loader has shuffle=False. Both use
                pin_memory=True for faster data transfer to GPU.
        """
        train, test = random_split(self.dataset, self.splits)
        train_loader = DataLoader(train, self.batch_size, shuffle=True, pin_memory=True)
        test_loader = DataLoader(test, self.batch_size, shuffle=False, pin_memory=True)
        return train_loader, test_loader

`init(dataset=None, splits=[0.8, 0.2], batch_size=16)` ¶

Initialize the DataLoader factory.

Parameters:

Name	Type	Description	Default
`dataset`	`Dataset`	PyTorch Dataset to split and load. Defaults to None.	`None`
`splits`	`list`	List of two floats representing train and test split proportions. Must sum to 1.0. Defaults to [0.8, 0.2].	`[0.8, 0.2]`
`batch_size`	`int`	Number of samples per batch. Defaults to 16.	`16`

Source code in deeplens\utils\dataset.py

def __init__(
        self, 
        dataset: Dataset = None,
        splits: list = [0.8, 0.2],
        batch_size: int = 16
    ) -> None:
    """Initialize the DataLoader factory.

    Args:
        dataset (Dataset, optional): PyTorch Dataset to split and load. Defaults to None.
        splits (list, optional): List of two floats representing train and test split
            proportions. Must sum to 1.0. Defaults to [0.8, 0.2].
        batch_size (int, optional): Number of samples per batch. Defaults to 16.
    """
    self.dataset = dataset
    self.splits = splits
    self.batch_size = batch_size

`deeplens.utils.tools` ¶

`get_device(device='auto')` ¶

Utility to set up the torch device. If 'auto', it selects the most appropriate device in your machine.

It can be set manually to 'mps', 'cuda', or 'cpu', but 'auto' is recommended.

Parameters:

Name	Type	Description	Default
`device`	`str`	The device in which the given process will be allocated. Defaults to 'auto'.	`'auto'`

Returns:

Type	Description
`device`	torch.device: The selected device.

Source code in deeplens\utils\tools.py

def get_device(device: str = "auto") -> torch.device:
    """Utility to set up the torch device. If 'auto', it selects
    the most appropriate device in your machine. 

    It can be set manually to 'mps', 'cuda', or 'cpu', but 'auto' is 
    recommended.

    Args:
        device: The device in which the given process will be allocated. 
            Defaults to 'auto'.

    Returns:
        torch.device: The selected device.
    """
    if device == "auto":
        return torch.device(
            "cuda" if torch.cuda.is_available() 
            else "mps" if torch.backends.mps.is_available()
            else "cpu"
        )
    return torch.device(device)

Utils API¶

deeplens.utils.analysis ¶

AnalysisUtils ¶

__init__(hf_model=None, sae_model=None, sae_config=None, layer=None, cache_dir='cache') ¶

generate_logit_heatmap(logits, save_name=None, k=None) ¶

get_most_active_features(sentences, target=None, k=None, case_sensitive=True, return_values=False) ¶

get_top_k_tokens(logits, k=10, to_dataframe=True, verbose=False) ¶

plot_topk_distribution(logits, k=20, position=0, save_name=None, use_softmax=True, title=None) ¶

deeplens.utils.dataset ¶

ActivationsDataset ¶

__getitem__(idx) ¶

__init__(activations) ¶

__len__() ¶

ActivationsDatasetBuilder ¶

__init__(activations=None, splits=[0.8, 0.2], batch_size=16, norm=True) ¶

get_dataloaders(ddp=False) ¶

normalize() ¶

set_tensor_dataset() ¶

AudioDatasetBuilder ¶

__getitem__(index) ¶

__init__(audio_dir=None, annotations_file=None, target_sample_rate=22050, num_samples=22050, device='auto', transformation_args={'n_fft': 1024, 'hop_length': 512, 'n_mels': 64}) ¶

__len__() ¶

GetDataLoaders ¶

__init__(dataset=None, splits=[0.8, 0.2], batch_size=16) ¶

deeplens.utils.tools ¶

get_device(device='auto') ¶

`deeplens.utils.analysis` ¶

`AnalysisUtils` ¶

`init(hf_model=None, sae_model=None, sae_config=None, layer=None, cache_dir='cache')` ¶

`generate_logit_heatmap(logits, save_name=None, k=None)` ¶

`get_most_active_features(sentences, target=None, k=None, case_sensitive=True, return_values=False)` ¶

`get_top_k_tokens(logits, k=10, to_dataframe=True, verbose=False)` ¶

`plot_topk_distribution(logits, k=20, position=0, save_name=None, use_softmax=True, title=None)` ¶

`deeplens.utils.dataset` ¶

`ActivationsDataset` ¶

`getitem(idx)` ¶

`init(activations)` ¶

`len()` ¶

`ActivationsDatasetBuilder` ¶

`init(activations=None, splits=[0.8, 0.2], batch_size=16, norm=True)` ¶

`get_dataloaders(ddp=False)` ¶

`normalize()` ¶

`set_tensor_dataset()` ¶

`AudioDatasetBuilder` ¶

`getitem(index)` ¶

`init(audio_dir=None, annotations_file=None, target_sample_rate=22050, num_samples=22050, device='auto', transformation_args={'n_fft': 1024, 'hop_length': 512, 'n_mels': 64})` ¶

`len()` ¶

`GetDataLoaders` ¶

`init(dataset=None, splits=[0.8, 0.2], batch_size=16)` ¶

`deeplens.utils.tools` ¶

`get_device(device='auto')` ¶