Datasets

Datasets are an extension point for declaring where prompts come from, with the same shape as an operation’s methods: a config whose type names an implementation in a registry, and a dataclass extending DatasetConfig for the settings that implementation takes. A dataset yields plain dict rows. entoscope ships no implementation: you register your own and DatasetFactory resolves it by type.

import json
from dataclasses import dataclass
from pathlib import Path

from entoscope.dataset import AbstractDataset, DatasetConfig, DatasetFactory, Row


@dataclass(kw_only=True)
class JsonlDatasetConfig(DatasetConfig):
    path: str
    limit: int | None = None


@DatasetFactory.register("jsonl")
class JsonlDataset(AbstractDataset[JsonlDatasetConfig]):
    def load(self) -> list[Row]:
        lines = Path(self.config.path).read_text().splitlines()
        return [json.loads(line) for line in lines][: self.config.limit]


Path("prompts.jsonl").write_text('{"prompt": "The capital of France is"}\n{"prompt": "The Eiffel Tower is in"}\n')

dataset = DatasetFactory.create(JsonlDatasetConfig(type="jsonl", path="prompts.jsonl", limit=10)).prepare()
print(len(dataset), dataset[0])  # 2 {'prompt': 'The capital of France is'}

In a config file:

my_data:
  type: jsonl
  path: ${DATA_DIR}/prompts.jsonl
  limit: 10

from entoscope import ConfigFactory

config = ConfigFactory[DatasetConfig]("configs.yaml").create("my_data")  # a JsonlDatasetConfig

A dataset’s type argument names its config class, as a method’s does: an entry is read strictly into the class the dataset its type names declares, so a missing setting or a key the class does not have fails as the file is read, and config.json reads back into the same class. A dataset with no settings of its own takes DatasetConfig itself, and its entry may be just its type.

load is called once, lazily, by prepare(); len, iteration and indexing then read the cached rows. Operations do not consume datasets yet: LensInput, ProbeInput, PatchingInput and CircuitInput take their prompts directly.

Registering a dataset follows the same rules as registering a method; see Registration.