Datasets¶

Functions and classes for sharing activation datasets via HuggingFace.

See the HuggingFace Datasets guide for a walkthrough.

Extraction¶

lmprobe.unified_cache.UnifiedCache ¶

Extracts activations and perplexity in a single forward pass.

This class provides efficient extraction by capturing both layer activations and lm_head logits in a single nnsight trace, then computing perplexity features from the logits.

Parameters:

Name	Type	Description	Default
`model`	`str`	HuggingFace model ID or local path.	required
`layers`	`int \| list[int] \| str`	Which layers to extract activations from: - int: Single layer (negative indexing supported) - list[int]: Multiple layers - "all": All layers (default) - "middle": Middle third of layers - "last": Last layer only	`"all"`
`compute_perplexity`	`bool`	Whether to also capture logits and compute perplexity features.	`True`
`device`	`str`	Device for model inference.	`"auto"`
`remote`	`bool`	Use nnsight remote execution (requires NDIF_API_KEY).	`False`
`batch_size`	`int`	Number of prompts to process at once.	`8`
`cache_pooled`	`bool`	If True (default), apply pooling BEFORE caching and store only the pooled activations. This reduces disk usage by ~100x (storing only (1, hidden_dim) per layer instead of (1, seq_len, hidden_dim)). Set to False only if you need to re-pool cached activations with different strategies. WARNING: When cache_pooled=True, you must use the same pooling strategy for all operations. The cached data cannot be re-pooled with a different strategy.	`True`
`pooling`	`str`	Pooling strategy to use when cache_pooled=True. Options: - "last_token": Use the last non-padding token (default) - "first_token": Use the first token - "mean": Mean of all non-padding tokens	`"last_token"`
`backend`	`str`	Extraction backend: "local" (default) or "nnsight".	`"local"`
`dtype`	`str or None`	Model dtype as a string: "float32", "float16", or "bfloat16". If None, defaults to float32 for local backend.	`None`

Examples:

>>> # Memory-efficient caching (recommended for large models)
>>> cache = UnifiedCache(
...     model="meta-llama/Llama-3.1-70B",
...     layers="all",
...     cache_pooled=True,      # Store only pooled activations
...     pooling="last_token",   # ~100x less disk space
... )
>>> stats = cache.warmup(prompts)
>>> print(f"Extracted {stats.activations_extracted} prompts")

model `property` ¶

model: Any

Get the loaded model, loading if necessary.

layer_indices `property` ¶

layer_indices: list[int]

Get resolved layer indices.

warmup ¶

warmup(prompts: list[str], remote: bool | None = None, max_retries: int | None = None, verify_cache: bool = True) -> WarmupStats

Extract and cache activations/perplexity for prompts.

This method checks the cache, identifies what needs extraction, and performs minimal forward passes to fill the cache.

Parameters:

Name	Type	Description	Default
`prompts`	`list[str]`	Text prompts to warm up the cache for.	required
`remote`	`bool \| None`	Override the instance-level remote setting.	`None`
`max_retries`	`int \| None`	Maximum number of retry attempts per batch for transient errors. Defaults to 3 for remote extraction, 0 for local.	`None`
`verify_cache`	`bool`	If True (default), read safetensors headers to verify that cached entries have the correct layers. If False, trust the cache key existence and skip header verification. Use False when you know the cache was populated with the same layer config (e.g. consecutive runs of the same extraction script) to avoid O(cached) S3 partial GETs.	`True`

Returns:

Type	Description
`WarmupStats`	Statistics about the warmup operation.

get_activations ¶

get_activations(prompts: list[str], remote: bool | None = None) -> tuple[torch.Tensor, torch.Tensor | None]

Get activations for prompts (from cache or via extraction).

Parameters:

Name	Type	Description	Default
`prompts`	`list[str]`	Text prompts.	required
`remote`	`bool \| None`	Override the instance-level remote setting.	`None`

Returns:

Type	Description
`tuple[Tensor, Tensor \| None]`	If cache_pooled=False: (activations, attention_mask) where activations has shape (batch, seq_len, n_layers * hidden_dim). If cache_pooled=True: (activations, None) where activations has shape (batch, n_layers * hidden_dim). Already pooled, no mask needed.

get_perplexity ¶

get_perplexity(prompts: list[str], remote: bool | None = None) -> np.ndarray

Get perplexity features for prompts (from cache or via extraction).

Parameters:

Name	Type	Description	Default
`prompts`	`list[str]`	Text prompts.	required
`remote`	`bool \| None`	Override the instance-level remote setting.	`None`

Returns:

Type	Description
`ndarray`	Perplexity features with shape (n_prompts, 3).

get_logits ¶

get_logits(prompts: list[str], remote: bool | None = None) -> CachedLogits

Get cached logits for prompts (from cache or via extraction).

Parameters:

Name	Type	Description	Default
`prompts`	`list[str]`	Text prompts.	required
`remote`	`bool \| None`	Override the instance-level remote setting.	`None`

Returns:

Type	Description
`CachedLogits`	Cached logits with values and optional indices tensors.

compute_logits_from_cache ¶

compute_logits_from_cache(prompts: list[str], positions: str = 'last', top_k: int | None = None, device: str | None = None, batch_size: int = 64) -> int

Compute logits from cached last-layer activations without a forward pass.

Downloads only the model's final norm and lm_head weights, then computes logits = norm(hidden_state) @ lm_head.T from the cached last-layer activations.

Parameters:

Name	Type	Description	Default
`prompts`	`list[str]`	Prompts whose last-layer activations are already cached.	required
`positions`	`str`	Which positions to compute logits for. Currently only "last" is supported.	`'last'`
`top_k`	`int \| None`	If set, store only top-k logit values and indices.	`None`
`device`	`str \| None`	Device for computation. Defaults to "cpu".	`None`
`batch_size`	`int`	Number of prompts to process at once for the matmul.	`64`

Returns:

Type	Description
`int`	Number of prompts for which logits were newly computed.

Raises:

Type	Description
`ValueError`	If the last layer is not in `self.layer_indices` or if `positions` is not "last".

compute_perplexity_from_cache ¶

compute_perplexity_from_cache(prompts: list[str], device: str | None = None, batch_size: int = 16) -> int

Compute perplexity from cached last-layer activations.

Reconstructs full-sequence logits as norm(hidden) @ lm_head.T, computes cross-entropy loss against input_ids (re-tokenized), and caches the perplexity stats (mean, min, max).

Full-sequence logits are NOT saved to disk (too large).

Parameters:

Name	Type	Description	Default
`prompts`	`list[str]`	Prompts whose last-layer raw activations are already cached.	required
`device`	`str \| None`	Device for computation. Defaults to "cpu".	`None`
`batch_size`	`int`	Number of prompts to process per batch.	`16`

Returns:

Type	Description
`int`	Number of prompts for which perplexity was newly computed.

Raises:

Type	Description
`ValueError`	If the last layer is not in `self.layer_indices`.

lmprobe.unified_cache.WarmupStats `dataclass` ¶

Statistics from a warmup operation.

cache_hit_rate `property` ¶

cache_hit_rate: float

Fraction of prompts that had activations cached.

Loading¶

load_activations(dataset: str, *, prompts: list[str] | None = None, layers: list[int] | None = None, pooling: str = 'last_token', token: str | None = None, as_dict: bool = True, return_labels: bool = False, show_progress: bool = True) -> dict[int, np.ndarray] | np.ndarray | tuple

Load pooled activations from a HuggingFace activation dataset.

Convenience function that downloads needed shards and returns structured activation arrays — no probe training required.

Parameters:

Name	Type	Description	Default
`dataset`	`str`	HuggingFace Dataset repo ID.	required
`prompts`	`list[str] \| None`	Prompts to load. None loads all prompts in the dataset.	`None`
`layers`	`list[int] \| None`	Layer indices to load. None loads all available layers.	`None`
`pooling`	`str`	Pooling strategy (default `"last_token"`).	`'last_token'`
`token`	`str \| None`	HuggingFace API token.	`None`
`as_dict`	`bool`	If True (default), return `{layer: ndarray(n_prompts, hidden_dim)}`. If False, return `ndarray(n_prompts, n_layers, hidden_dim)`.	`True`
`return_labels`	`bool`	If True, also return labels from the dataset's Parquet index as a second element: `(activations, labels)`. Labels are a numpy array of ints (or None if the dataset has no `label` column).	`False`
`show_progress`	`bool`	If True (default), display tqdm progress bars for downloads and activation loading.	`True`

Returns:

Type	Description
`dict[int, ndarray] \| ndarray \| tuple`	Activation arrays keyed by layer index, or a single stacked array. If `return_labels=True`, returns `(activations, labels)` where labels is `np.ndarray \| None`.

Publishing¶

push_dataset(repo_id: str, model_name: str, prompts: list[str], *, labels: list[int | str | None] | None = None, metadata: list[dict] | None = None, tensors: list[str] | None = None, shard_max_bytes: int = DEFAULT_SHARD_BYTES, private: bool = False, exist_ok: bool = False, skip_missing: bool = True, description: str | None = None, license: str = 'cc-by-4.0', token: str | None = None, num_workers: int | None = None, commit_batch_size: int | None = None, stream: bool = False, stream_batch_size: int = 10, preload: str = 'none', preload_workers: int = 8, shuffle: bool = True) -> str

Push cached activations to a HuggingFace Dataset repo.

Parameters:

Name	Type	Description	Default
`repo_id`	`str`	HuggingFace repo ID (e.g. `"username/my-activations"`).	required
`model_name`	`str`	The model whose activations are cached.	required
`prompts`	`list[str]`	Prompts to push (must have cached activations).	required
`labels`	`list[int \| str \| None] \| None`	Per-prompt labels, positionally aligned with prompts.	`None`
`metadata`	`list[dict] \| None`	Per-prompt metadata dicts, positionally aligned with prompts. All dicts must have the same keys. Values appear as extra columns in the Parquet index.	`None`
`tensors`	`list[str] \| None`	Filter: only push these tensor types (`["hidden_layers"]`, `["logits_topk"]`). None pushes all available types.	`None`
`shard_max_bytes`	`int`	Max bytes per shard file. Default 1 GB.	`DEFAULT_SHARD_BYTES`
`private`	`bool`	Create a private repository.	`False`
`exist_ok`	`bool`	If False (default), raise if the repo already exists.	`False`
`skip_missing`	`bool`	If True (default), skip prompts missing from cache.	`True`
`description`	`str \| None`	Description for the auto-generated README.	`None`
`license`	`str`	License identifier for the dataset card. Default `"cc-by-4.0"`.	`'cc-by-4.0'`
`token`	`str \| None`	HuggingFace API token.	`None`
`num_workers`	`int \| None`	Number of workers for parallel file uploads. Passed directly to `upload_large_folder`. `None` (default) uses the `huggingface_hub` default (currently 16 workers). Ignored when `stream=True`.	`None`
`commit_batch_size`	`int \| None`	Maximum number of files per commit during upload. On unreliable connections, setting this to a small value (e.g. `1`) ensures progress is committed frequently, so interrupted uploads lose less work on restart. `None` (default) uses the `huggingface_hub` default scale. Ignored when `stream=True`.	`None`
`stream`	`bool`	If True, upload shards in batches after writing them, then delete local copies. This reduces peak disk usage from the full dataset size to `stream_batch_size` shards (~N GB). Resumable via a shard completion manifest. Default False.	`False`
`stream_batch_size`	`int`	Number of shards to buffer before uploading as a single commit in streaming mode. Higher values improve throughput (fewer commits, parallel LFS uploads) at the cost of more peak disk. Default 10. Ignored when `stream=False`.	`10`
`preload`	`str`	Strategy for preloading cached tensors during consolidation. `"none"` (default) reads each prompt sequentially per layer/shard. `"per_layer"` preloads one layer across all prompts in parallel before writing shards — best trade-off for most cases. `"full"` preloads all layers for all prompts into memory — fastest but requires enough RAM to hold the entire dataset.	`'none'`
`preload_workers`	`int`	Number of parallel threads for preloading cache reads. Only used when `preload` is not `"none"`. Default 8.	`8`
`shuffle`	`bool`	If True (default), deterministically shuffle prompts across shards using a seed derived from `repo_id`. If False, preserve the input prompt order in the output shards.	`True`

Returns:

Type	Description
`str`	URL of the dataset.

pull_dataset(repo_id: str, *, tensors: list[str] | None = None, layers: list[int] | None = None, target_prompts: list[str] | None = None, overwrite: bool = False, token: str | None = None, materialize: bool = False, num_workers: int = 0, show_progress: bool = False) -> int

Pull activations from a HuggingFace Dataset repo into local cache.

By default (materialize=False), downloads shard files and builds a shard registry so activations can be served lazily on-the-fly without unpacking per-prompt files. Set materialize=True to unpack per-prompt safetensors files (the pre-0.8 behavior).

Parameters:

Name	Type	Description	Default
`repo_id`	`str`	HuggingFace repo ID.	required
`tensors`	`list[str] \| None`	Only pull these tensor types (`["hidden_layers"]`, `["logits_topk"]`). None pulls all.	`None`
`layers`	`list[int] \| None`	Only download these hidden layers (e.g. `[0, 1]`). Requires per-layer layout (v1.1). On v1.0 co-located datasets, this parameter is ignored with a warning. None downloads all layers.	`None`
`target_prompts`	`list[str] \| None`	Only pull these prompts. None pulls all.	`None`
`overwrite`	`bool`	If False (default), skip prompts already in local cache.	`False`
`token`	`str \| None`	HuggingFace API token.	`None`
`materialize`	`bool`	If True, unpack per-prompt safetensors files into the cache (old behavior). If False (default), only build the shard registry for lazy on-the-fly loading.	`False`
`num_workers`	`int`	Number of parallel workers for materialization. 0 (default) runs in the main process. Only used when `materialize=True`.	`0`
`show_progress`	`bool`	If True, display a tqdm progress bar for shard downloads.	`False`

Returns:

Type	Description
`int`	Number of prompts available (lazy) or unpacked (materialize).

load_activation_dataset(repo_id: str, *, tensors: list[str] | None = None, layers: list[int] | None = None, token: str | None = None) -> tuple[dict[str, torch.Tensor], dict]

Load tensors directly from a HuggingFace Dataset repo.

Downloads and concatenates shards, returns raw tensors and metadata. No local cache interaction.

Parameters:

Name	Type	Description	Default
`repo_id`	`str`	HuggingFace repo ID.	required
`tensors`	`list[str] \| None`	Only load these tensor types (`["hidden_layers"]`, `["logits_topk"]`). None loads all.	`None`
`layers`	`list[int] \| None`	Only download these hidden layers. Requires per-layer layout (v1.1). On v1.0 co-located datasets, ignored with a warning. None downloads all layers.	`None`
`token`	`str \| None`	HuggingFace API token.	`None`

Returns:

Type	Description
`tuple[dict[str, Tensor], dict]`	(tensors_by_key, lmprobe_info) where tensors_by_key maps safetensors keys (e.g. `"hidden.layer_0"`) to concatenated tensors.

fetch_dataset_metadata(repo_id: str, *, token: str | None = None) -> DatasetMetadata

Fetch lightweight metadata from a remote activation dataset.

Downloads the Parquet index (which contains embedded metadata in v2 format) — no tensor data is transferred. Falls back to downloading lmprobe_info.json for v1 datasets.

Parameters:

Name	Type	Description	Default
`repo_id`	`str`	HuggingFace Dataset repo ID (e.g. `"user/my-activations"`).	required
`token`	`str \| None`	HuggingFace API token.	`None`

Returns:

Type	Description
`DatasetMetadata`	Parsed metadata including model name, available layers, prompt list, and tensor descriptors.

migrate_dataset(repo_id: str, *, shard_max_bytes: int = DEFAULT_SHARD_BYTES, token: str | None = None, private: bool | None = None, dry_run: bool = False) -> str

Migrate a full_sequence dataset to use last-token shard splitting.

Downloads the existing dataset from HuggingFace, re-shards the hidden layer tensors so that last-token vectors land in dedicated small shards, rebuilds the parquet index with per-token shard mapping arrays (token_shard_ids, token_shard_offsets), and re-uploads.

Processing is done one layer at a time to limit memory usage.

Parameters:

Name	Type	Description	Default
`repo_id`	`str`	HuggingFace repo ID of the dataset to migrate.	required
`shard_max_bytes`	`int`	Max bytes per shard file for the new sharding. Default 1 GB.	`DEFAULT_SHARD_BYTES`
`token`	`str \| None`	HuggingFace API token.	`None`
`private`	`bool \| None`	If set, update the repo visibility. None keeps the current setting.	`None`
`dry_run`	`bool`	If True, download and compute the new plan but do not upload.	`False`

Returns:

Type	Description
`str`	URL of the migrated dataset (or a summary string if dry_run).

upgrade_dataset_format(repo_id: str, *, token: str | None = None, remove_json: bool = True, dry_run: bool = False) -> str

Upgrade a v1.x dataset to v2 format (metadata in Parquet schema).

This is a lightweight, metadata-only operation — no tensor files are modified. It reads the existing lmprobe_info.json and embeds its contents into the Parquet index's schema metadata under lmprobe: prefixed keys.

Parameters:

Name	Type	Description	Default
`repo_id`	`str`	HuggingFace repo ID of the dataset to upgrade.	required
`token`	`str \| None`	HuggingFace API token.	`None`
`remove_json`	`bool`	If True (default), delete `lmprobe_info.json` from the repo after embedding metadata in Parquet.	`True`
`dry_run`	`bool`	If True, download and validate but do not upload changes.	`False`

Returns:

Type	Description
`str`	URL of the upgraded dataset (or a summary string if dry_run).

Metadata from a remote activation dataset.

Returned by :func:fetch_dataset_metadata. Contains just enough information to resolve layers, validate model compatibility, and check prompt availability without downloading any tensor data.

Datasets¶

Extraction¶

lmprobe.unified_cache.UnifiedCache ¶

model property ¶

layer_indices property ¶

warmup ¶

get_activations ¶

get_perplexity ¶

get_logits ¶

compute_logits_from_cache ¶

compute_perplexity_from_cache ¶

lmprobe.unified_cache.WarmupStats dataclass ¶

cache_hit_rate property ¶

Loading¶

lmprobe.sharing.load_activations ¶

Publishing¶

lmprobe.sharing.push_dataset ¶

lmprobe.sharing.pull_dataset ¶

lmprobe.sharing.load_activation_dataset ¶

lmprobe.sharing.fetch_dataset_metadata ¶

lmprobe.sharing.migrate_dataset ¶

lmprobe.sharing.upgrade_dataset_format ¶

lmprobe.sharing.DatasetMetadata dataclass ¶

model `property` ¶

layer_indices `property` ¶

lmprobe.unified_cache.WarmupStats `dataclass` ¶

cache_hit_rate `property` ¶

lmprobe.sharing.DatasetMetadata `dataclass` ¶