Annotator¶

llm_annotator.annotator ¶

Annotator `dataclass` ¶

Annotator(
    client: Client,
    batch_size: int = 256,
    num_proc: int | None = DEFAULT_CPU_COUNT,
    verbose: bool = False,
)

Sensible base class for LLM-based dataset annotation.

This class provides a framework for annotating datasets using LLMs via a pluggable Client. It handles dataset loading, processing, and output generation with support for batching and uploading to the Hugging Face Hub.

The Annotator class has four public entry points:

prepare_data. Apply prompt templates, sorting, and caching without running inference. Backs-up prepared artifacts to Hugging Face Hub if hub_id is provided.
run_annotation. Run inference only, using prepared data returned by prepare_data or loaded from a local path or Hub repo.
annotate_dataset. Convenience wrapper that calls prepare_data and then run_annotation in one call.
generate_dataset. Generate a new dataset from scratch by calling annotate_dataset over a synthetic prompt dataset.

For large-scale annotation jobs, consider using VLLMQueueAnnotator, which distributes inference across a pool of vLLM servers.

Parameters:

Name	Type	Description	Default
`client`	`Client`	An initialised `Client` instance that performs the actual generation.	required
`batch_size`	`int`	Number of samples per inference batch. It depends on the client and its settings which batching is actually used. Batch size here is mostly intended for progress reporting. The client may split the given batch into smaller sub-batches if needed.	`256`
`num_proc`	`int \| None`	Number of processes for dataset preprocessing.	`DEFAULT_CPU_COUNT`
`verbose`	`bool`	Whether to print progress information.	`False`

Examples:

Basic usage with an OpenAI client:

from llm_annotator import Annotator, OpenAIClient
client = OpenAIClient(model="gpt-4o-mini")
with Annotator(client=client) as anno:
    ds = anno.annotate_dataset(
        output_dir="outputs/data",
        prompt_template="Process: {text}",
        dataset_name="my-dataset",
    )

Usage with vLLM offline client:

from llm_annotator import Annotator, VLLMOfflineClient
client = VLLMOfflineClient(
    model="meta-llama/Llama-3.2-3B-Instruct",
    max_model_len=4096,
)
try:
    ds = Annotator(client=client).annotate_dataset(
        output_dir="outputs/data",
        prompt_template="Process: {text}",
        dataset_name="my-dataset",
    )
finally:
    client.destroy()

__post_init__ ¶

__post_init__() -> None

Initialize the logger for annotator runtime messages.

View source on GitHub: src/llm_annotator/annotator.py lines 185–187

enter ¶

__enter__() -> 'Annotator'

Enter the context manager, returning the annotator instance.

View source on GitHub: src/llm_annotator/annotator.py lines 189–191

exit ¶

__exit__(exc_type: Any, exc: Any, tb: Any) -> None

Exit the context manager and free all client resources.

View source on GitHub: src/llm_annotator/annotator.py lines 193–195

destroy ¶

destroy() -> None

Clean up all resources used by the underlying client.

View source on GitHub: src/llm_annotator/annotator.py lines 197–199

prepare_data ¶

prepare_data(
    output_dir: str | Path,
    prompt_template: str,
    *,
    dataset_name: str | None = None,
    dataset: Dataset | None = None,
    dataset_config: str | None = None,
    data_dir: str | None = None,
    data_files: str
    | list[str]
    | dict[str, str | list[str]]
    | None = None,
    dataset_split: str | None = None,
    max_num_samples: int | None = None,
    shuffle_seed: int | None = None,
    preprocess_fn: Callable | None = None,
    prompt_field_swapper: dict[str, str] | None = None,
    idx_column: str = "idx",
    task_prefix: str = "",
    sort_by_length: bool
    | Literal["shortest_first", "longest_first"] = False,
    system_message: str | None = None,
    hub_id: str | None = None,
    keep_columns: str | Iterable[str] | bool | None = None,
    force_data_preparation: bool = False,
) -> tuple[Dataset, Path | None, str | None]

Prepare input data for annotation without running generation.

The method reuses local prepared data first, then optionally restores prepared data from Hugging Face Hub, and finally falls back to building the prepared dataset from source.

Only the columns required for inference are retained in the cached artifact: idx_column and {task_prefix}messages. Pass keep_columns to preserve additional source columns (e.g. those needed by run_annotation's keep_columns argument).

Parameters:

Name	Type	Description	Default
`output_dir`	`str \| Path`	Directory where prepared artifacts are stored.	required
`prompt_template`	`str`	Prompt template used to build chat messages.	required
`dataset_name`	`str \| None`	Name or path of the dataset to load.	`None`
`dataset`	`Dataset \| None`	Pre-loaded dataset to use instead of loading from name/path.	`None`
`dataset_config`	`str \| None`	Dataset configuration name (optional).	`None`
`data_dir`	`str \| None`	Data directory for local datasets (optional).	`None`
`data_files`	`str \| list[str] \| dict[str, str \| list[str]] \| None`	Specific file(s) for local datasets (optional).	`None`
`dataset_split`	`str \| None`	Specific split to load (optional).	`None`
`max_num_samples`	`int \| None`	Maximum number of samples to prepare.	`None`
`shuffle_seed`	`int \| None`	Seed for dataset shuffling.	`None`
`preprocess_fn`	`Callable \| None`	Optional function to preprocess the dataset after loading and before ap plying the prompt template.	`None`
`prompt_field_swapper`	`dict[str, str] \| None`	Optional mapping to replace template fields.	`None`
`idx_column`	`str`	Column name used as unique identifier. Must not exist in the input dataset.	`'idx'`
`task_prefix`	`str`	Prefix for internal columns and artifact names.	`''`
`sort_by_length`	`bool \| Literal['shortest_first', 'longest_first']`	Whether to sort prompts by length.	`False`
`system_message`	`str \| None`	Optional system message for chat prompts.	`None`
`hub_id`	`str \| None`	Optional Hugging Face dataset ID used for both prepared-data backup and restore. Will be stored in the PREPARED_DS_BRANCH_SUFF branch.	`None`
`keep_columns`	`str \| Iterable[str] \| bool \| None`	Source columns to retain in the cached artifact in addition to the essential `idx_column` and messages column. `True` keeps all columns (logs a size warning). `None` or an empty collection keeps only the essential columns.	`None`
`force_data_preparation`	`bool`	Whether to rebuild prepared data even when local or Hub artifacts already exist.	`False`

Returns:

Type	Description
`Dataset`	Tuple of prepared dataset, local prepared-data path when available,
`Path \| None`	and Hugging Face dataset ID when available.

View source on GitHub: src/llm_annotator/annotator.py lines 840–1019

run_annotation ¶

run_annotation(
    output_dir: str | Path,
    prompt_template: str | None = None,
    *,
    prepared_dataset: Dataset | None = None,
    prepared_data_path: str | Path | None = None,
    hub_id: str | None = None,
    overwrite: bool = False,
    dataset_split: str | None = None,
    dataset_config: str | None = None,
    keep_columns: str | Iterable[str] | bool | None = None,
    options: ProviderRuntimeOptions | None = None,
    gen_kwargs: dict[str, Any] | None = None,
    output_schema: str | dict[str, Any] | None = None,
    idx_column: str = "idx",
    upload_every_n_samples: int | None = 10000,
    max_samples_per_output_file: int = 1000,
    task_prefix: str = "",
    validate_fn: Callable | None = None,
    postprocess_fn: Callable | None = None,
    num_retries_invalid: int = 5,
    system_message: str | None = None,
    keep_idx_column: bool = False,
    max_consecutive_failed_batches: int = 10,
) -> Dataset

Run model generation on already prepared annotation inputs.

Parameters:

Name	Type	Description	Default
`output_dir`	`str \| Path`	Directory where annotation output is written.	required
`prompt_template`	`str \| None`	Prompt template used for warm-up metadata. Optional because the prepared dataset already carries the rendered messages; when given, its static prefix is used to prime the prefix cache.	`None`
`prepared_dataset`	`Dataset \| None`	Pre-prepared dataset with messages column.	`None`
`prepared_data_path`	`str \| Path \| None`	Local path to prepared data on disk.	`None`
`hub_id`	`str \| None`	Hugging Face dataset ID used for prepared-data cache and JSONL progress backup.	`None`
`overwrite`	`bool`	Whether to overwrite existing output directory EXCEPT for the prepared data cache (which is preserved to allow resuming). If you want to overwrite the prepared data cache, delete it manually or set `force_data_preparation=True` in `prepare_data`.	`False`
`dataset_split`	`str \| None`	Dataset split used for skip filtering.	`None`
`dataset_config`	`str \| None`	Dataset config used for skip filtering.	`None`
`keep_columns`	`str \| Iterable[str] \| bool \| None`	Columns to keep in output. `True` for all.	`None`
`options`	`ProviderRuntimeOptions \| None`	Runtime options passed to the client.	`None`
`gen_kwargs`	`dict[str, Any] \| None`	Extra request parameters merged over `options`, for anything the options dataclass does not name.	`None`
`output_schema`	`str \| dict[str, Any] \| None`	Convenience JSON schema input. When provided, it is injected into `options.json_schema`.	`None`
`idx_column`	`str`	Column name used as unique identifier.	`'idx'`
`upload_every_n_samples`	`int \| None`	Upload to Hub every N samples.	`10000`
`max_samples_per_output_file`	`int`	Maximum samples per output file.	`1000`
`task_prefix`	`str`	Prefix for internal columns and file names.	`''`
`validate_fn`	`Callable \| None`	Optional custom validation function.	`None`
`postprocess_fn`	`Callable \| None`	Optional postprocessing function that takes in a sample and must return a dict.	`None`
`num_retries_invalid`	`int`	Number of retries for invalid outputs.	`5`
`system_message`	`str \| None`	Optional system message for chat prompts.	`None`
`keep_idx_column`	`bool`	Whether to keep idx column in final dataset.	`False`
`max_consecutive_failed_batches`	`int`	Abort the run once this many batches in a row come back with every sample errored (e.g. a vLLM server that died mid-run), instead of continuing to dispatch batches against a backend that isn't responding. Set to 0 to disable.	`10`

Returns:

Type	Description
`Dataset`	Final concatenated annotation dataset.

Raises:

Type	Description
`ValueError`	If no prepared data source can be resolved.
`TooManyConsecutiveFailedBatchesError`	If `max_consecutive_failed_batches` consecutive batches fail entirely.

View source on GitHub: src/llm_annotator/annotator.py lines 1021–1368

annotate_dataset ¶

annotate_dataset(
    output_dir: str | Path,
    prompt_template: str | None = None,
    *,
    full_prompt_template: str | None = None,
    dataset_name: str | None = None,
    dataset: Dataset | None = None,
    dataset_config: str | None = None,
    data_dir: str | None = None,
    data_files: str
    | list[str]
    | dict[str, str | list[str]]
    | None = None,
    dataset_split: str | None = None,
    max_num_samples: int | None = None,
    shuffle_seed: int | None = None,
    preprocess_fn: Callable | None = None,
    prompt_field_swapper: dict[str, str] | None = None,
    idx_column: str = "idx",
    task_prefix: str = "",
    sort_by_length: bool
    | Literal["shortest_first", "longest_first"] = False,
    system_message: str | None = None,
    hub_id: str | None = None,
    force_data_preparation: bool = False,
    overwrite: bool = False,
    keep_columns: str | Iterable[str] | bool | None = None,
    options: ProviderRuntimeOptions | None = None,
    gen_kwargs: dict[str, Any] | None = None,
    output_schema: str | dict[str, Any] | None = None,
    upload_every_n_samples: int | None = 10000,
    max_samples_per_output_file: int = 1000,
    validate_fn: Callable | None = None,
    postprocess_fn: Callable | None = None,
    num_retries_invalid: int = 5,
    keep_idx_column: bool = False,
    max_consecutive_failed_batches: int = 10,
) -> Dataset

Annotate an existing dataset in one call.

This is a convenience wrapper around prepare_data and run_annotation for callers that prefer a single entry point.

Parameters:

Name	Type	Description	Default
`output_dir`	`str \| Path`	Directory where annotation output is written.	required
`prompt_template`	`str \| None`	Prompt template with dataset fields. Defaults to `full_prompt_template` when provided.	`None`
`full_prompt_template`	`str \| None`	Backwards-compatible alias for `prompt_template`.	`None`
`dataset_name`	`str \| None`	Name or path of the dataset to load.	`None`
`dataset`	`Dataset \| None`	Pre-loaded dataset to annotate instead of loading one.	`None`
`dataset_config`	`str \| None`	Dataset configuration name.	`None`
`data_dir`	`str \| None`	Data directory for local datasets.	`None`
`data_files`	`str \| list[str] \| dict[str, str \| list[str]] \| None`	Specific file(s) for local datasets.	`None`
`dataset_split`	`str \| None`	Dataset split to load.	`None`
`max_num_samples`	`int \| None`	Maximum number of samples to annotate.	`None`
`shuffle_seed`	`int \| None`	Seed for dataset shuffling.	`None`
`preprocess_fn`	`Callable \| None`	Optional preprocessing callback.	`None`
`prompt_field_swapper`	`dict[str, str] \| None`	Optional mapping that renames prompt fields.	`None`
`idx_column`	`str`	Column name used as the stable sample identifier.	`'idx'`
`task_prefix`	`str`	Prefix for internal column names and output files.	`''`
`sort_by_length`	`bool \| Literal['shortest_first', 'longest_first']`	Whether to sort prompts by length.	`False`
`system_message`	`str \| None`	Optional system message for the chat prompt.	`None`
`hub_id`	`str \| None`	Optional Hub dataset ID for prepared-data cache and JSONL progress backup.	`None`
`force_data_preparation`	`bool`	Rebuild prepared data even if cached.	`False`
`overwrite`	`bool`	Whether to overwrite the output directory EXCEPT for the prepared data cache (which is preserved to allow resuming). If you want to overwrite the prepared data cache, delete it manually or set `force_data_preparation=True`.	`False`
`keep_columns`	`str \| Iterable[str] \| bool \| None`	Columns to keep in the final dataset.	`None`
`options`	`ProviderRuntimeOptions \| None`	Runtime options passed to the client.	`None`
`gen_kwargs`	`dict[str, Any] \| None`	Extra request parameters merged over `options`, for anything the options dataclass does not name.	`None`
`output_schema`	`str \| dict[str, Any] \| None`	Optional JSON schema for structured output.	`None`
`upload_every_n_samples`	`int \| None`	Upload checkpoint cadence.	`10000`
`max_samples_per_output_file`	`int`	Maximum samples per output file.	`1000`
`validate_fn`	`Callable \| None`	Optional validation callback.	`None`
`postprocess_fn`	`Callable \| None`	Optional postprocessing callback.	`None`
`num_retries_invalid`	`int`	Number of retries for invalid outputs.	`5`
`keep_idx_column`	`bool`	Whether to keep the index column in the result.	`False`
`max_consecutive_failed_batches`	`int`	Abort the run once this many batches in a row come back with every sample errored. Set to 0 to disable.	`10`

Returns:

Type	Description
`Dataset`	The concatenated annotation dataset.

Raises:

Type	Description
`TypeError`	If no prompt template is provided.
`TooManyConsecutiveFailedBatchesError`	If `max_consecutive_failed_batches` consecutive batches fail entirely.

View source on GitHub: src/llm_annotator/annotator.py lines 1370–1521

generate_dataset ¶

generate_dataset(
    output_dir: str | Path,
    prompts: str | Sequence[str],
    *,
    prompt_prefix: str | None = None,
    hub_id: str | None = None,
    force_data_preparation: bool = False,
    overwrite: bool = False,
    options: ProviderRuntimeOptions | None = None,
    gen_kwargs: dict[str, Any] | None = None,
    max_num_samples: int | None = None,
    output_schema: str | dict[str, Any] | None = None,
    idx_column: str = "idx",
    upload_every_n_samples: int | None = 10000,
    max_samples_per_output_file: int = 1000,
    task_prefix: str = "",
    validate_fn: Callable | None = None,
    postprocess_fn: Callable | None = None,
    num_retries_invalid: int = 5,
    keep_idx_column: bool = False,
    max_consecutive_failed_batches: int = 10,
) -> Dataset

Generate a new dataset from prompts.

Parameters:

Name	Type	Description	Default
`output_dir`	`str \| Path`	Directory where annotation output is written.	required
`prompts`	`str \| Sequence[str]`	A single prompt or a sequence of prompts.	required
`prompt_prefix`	`str \| None`	Optional shared prefix used for prefix caching.	`None`
`hub_id`	`str \| None`	Optional Hub dataset ID for prepared-data cache and JSONL progress backup.	`None`
`force_data_preparation`	`bool`	Rebuild prepared data even if cached.	`False`
`overwrite`	`bool`	Whether to overwrite the output directory EXCEPT for the prepared data cache (which is preserved to allow resuming). If you want to overwrite the prepared data cache, delete it manually or set `force_data_preparation=True`.	`False`
`options`	`ProviderRuntimeOptions \| None`	Runtime options passed to the client.	`None`
`gen_kwargs`	`dict[str, Any] \| None`	Extra request parameters merged over `options`, for anything the options dataclass does not name.	`None`
`max_num_samples`	`int \| None`	Number of times to repeat a single prompt.	`None`
`output_schema`	`str \| dict[str, Any] \| None`	Optional JSON schema for structured output.	`None`
`idx_column`	`str`	Column name used as the stable sample identifier.	`'idx'`
`upload_every_n_samples`	`int \| None`	Upload checkpoint cadence.	`10000`
`max_samples_per_output_file`	`int`	Maximum samples per output file.	`1000`
`task_prefix`	`str`	Prefix for internal column names and output files.	`''`
`validate_fn`	`Callable \| None`	Optional validation callback.	`None`
`postprocess_fn`	`Callable \| None`	Optional postprocessing callback.	`None`
`num_retries_invalid`	`int`	Number of retries for invalid outputs.	`5`
`keep_idx_column`	`bool`	Whether to keep the index column in the result.	`False`
`max_consecutive_failed_batches`	`int`	Abort the run once this many batches in a row come back with every sample errored. Set to 0 to disable.	`10`

Returns:

Type	Description
`Dataset`	The concatenated annotation dataset.

Raises:

Type	Description
`ValueError`	If no prompts are provided.
`TooManyConsecutiveFailedBatchesError`	If `max_consecutive_failed_batches` consecutive batches fail entirely.

View source on GitHub: src/llm_annotator/annotator.py lines 1523–1628

get_pfout_name ¶

get_pfout_name(
    *,
    process_pdout: Path,
    max_samples_per_output_file: int,
    processed_n_samples: int | None = None,
) -> Path

Generate the output file name based on configuration.

Creates appropriate file names for output files, handling both single-file and multi-file output modes.

Parameters:

Name	Type	Description	Default
`process_pdout`	`Path`	The output directory path.	required
`max_samples_per_output_file`	`int`	Maximum samples per output file (0 for unlimited).	required
`processed_n_samples`	`int \| None`	The number of samples processed so far.	`None`

Returns:

Type	Description
`Path`	Path object for the output file name.

View source on GitHub: src/llm_annotator/annotator.py lines 1862–1888

push_progress_to_hub ¶

push_progress_to_hub(
    dir_path: Path | str,
    hub_id: str | None = None,
    *,
    task_prefix: str = "",
) -> None

Upload the output directory to Hugging Face Hub.

Creates a dataset repository and uploads all annotation files, excluding cached input data. Uses a separate branch for uploads.

Parameters:

Name	Type	Description	Default
`dir_path`	`Path \| str`	Path to the directory containing annotation files.	required
`hub_id`	`str \| None`	Optional Hugging Face dataset ID to upload into.	`None`
`task_prefix`	`str`	String prefix to use for branch naming.	`''`

View source on GitHub: src/llm_annotator/annotator.py lines 1890–1932

VLLMQueueAnnotator `dataclass` ¶

VLLMQueueAnnotator(
    client: Client,
    batch_size: int = 256,
    num_proc: int | None = DEFAULT_CPU_COUNT,
    verbose: bool = False,
    *,
    clients: Sequence[Client[Any]],
    queue_size: int | None = None,
)

Bases: Annotator

Annotator that spreads one workload over several vLLM servers/clients.

Instead of walking the prepared dataset batch-by-batch through a single client it keeps a bounded queue of batches in flight over a pool of vLLM server clients, handing each batch to whichever server is free. The process can be simplified as:

add all clients to a queue;
for each batch:
- pop a client from the queue;
- send the batch to that client;

Everything else -- prompt templating, JSONL progress snapshots keyed by idx, resumption, Hub backups, and the final concatenation -- is inherited from Annotator, so all four public entry points behave exactly as documented there.

Because batches finish out of order, results are written in completion order and sorted by idx at the end (as they already are for the base annotator, whose JSONL files are concatenated and sorted in Annotator._post_annotate).

Parameters:

Name	Type	Description	Default
`clients`	`Sequence[Client[Any]]`	vLLM server clients used as the worker pool. Keyword-only. The first client doubles as `Annotator.client` for inherited helpers, so `client` is derived here rather than passed in.	required
`queue_size`	`int \| None`	Maximum number of batches in flight (dispatched but not yet written out). This bounds memory, not the amount of work: the full dataset is always annotated. `None` resolves to four batches per client, and any value below `len(clients)` is raised to it, since a smaller queue would leave servers idle. After initialisation the attribute always holds the resolved value.	`None`
`batch_size`	`int`	Maximum number of samples sent to a worker in one request.	`256`
`num_proc`	`int \| None`	Number of processes for dataset preprocessing.	`DEFAULT_CPU_COUNT`
`verbose`	`bool`	Whether to print progress information.	`False`

Raises:

Type	Description
`ValueError`	If no clients are given or `queue_size` is not positive.
`TypeError`	If a client is not a vLLM server client.

Examples:

from llm_annotator import VLLMOnlineClient, VLLMQueueAnnotator
clients = [
    VLLMOnlineClient(
        model="Qwen/Qwen3-8B", base_url=f"http://{host}:8000/v1"
    )
    for host in ("node01", "node02")
]
with VLLMQueueAnnotator(
    clients=clients, batch_size=64
) as anno:
    ds = anno.annotate_dataset(
        output_dir="outputs/data",
        prompt_template="Classify: {text}",
        dataset_name="my-dataset",
    )

__post_init__ ¶

__post_init__() -> None

Validate the pool, derive the defaults, and fill the client queue.

Raises:

Type	Description
`ValueError`	If no clients are given or `queue_size` is not positive.
`TypeError`	If a client is not a vLLM server client.

View source on GitHub: src/llm_annotator/annotator.py lines 2047–2085

set_queue_size ¶

set_queue_size(queue_size: int | None) -> None

Change how many batches are kept in flight.

Assigning to queue_size directly would break the invariant that it always holds a resolved value, since None and values below the pool size are only normalised on the way in. Use this instead when a pool is reused for another workload.

Parameters:

Name	Type	Description	Default
`queue_size`	`int \| None`	Requested number of batches in flight, or `None` to derive it from the pool size.	required

Raises:

Type	Description
`ValueError`	If `queue_size` is given but not positive.

View source on GitHub: src/llm_annotator/annotator.py lines 2119–2134

destroy ¶

destroy() -> None

Clean up the resources of every client in the pool. Since clients can only be VLLMOnlineClients, the impact is likely minimal: that class has no meaningful destroy of its own. It inherits OpenAIClient's, which only does batch-related cleanup, and vLLM does not support the OpenAI Batch API.

Every client is destroyed even if some of them raise; the first error is re-raised afterwards.

Raises:

Type	Description
`BaseException`	The first error raised by a client, if any.

View source on GitHub: src/llm_annotator/annotator.py lines 2136–2161

destroy_on_error ¶

destroy_on_error(
    func: Callable[..., Any],
) -> Callable[..., Any]

Decorate an Annotator method to clean up on any exception.

Calls destroy before re-raising. Catches BaseException (including KeyboardInterrupt and SystemExit) so resources are freed even on forced termination. The original exception is always re-raised after the cleanup attempt.

Should be used on methods that use the underlying client.

Parameters:

Name	Type	Description	Default
`func`	`Callable[..., Any]`	The instance method to wrap.	required

Returns:

Type	Description
`Callable[..., Any]`	The wrapped callable with automatic cleanup on failure.

View source on GitHub: src/llm_annotator/annotator.py lines 79–108

Annotator¶

llm_annotator.annotator ¶

Annotator dataclass ¶

__post_init__ ¶

__enter__ ¶

__exit__ ¶

destroy ¶

prepare_data ¶

run_annotation ¶

annotate_dataset ¶

generate_dataset ¶

get_pfout_name ¶

push_progress_to_hub ¶

VLLMQueueAnnotator dataclass ¶

__post_init__ ¶

set_queue_size ¶

destroy ¶

destroy_on_error ¶

Annotator `dataclass` ¶

enter ¶

exit ¶

VLLMQueueAnnotator `dataclass` ¶