Annotator¶

llm_annotator.annotator ¶

BOOKKEEPING_SUFFIXES `module-attribute` ¶

BOOKKEEPING_SUFFIXES = (
    "error",
    "error_type",
    "finish_reason",
    "messages",
    "num_tokens",
    "reasoning",
    "response",
    "valid",
    "valid_fields",
)

Column names, after the task prefix, that the annotator writes itself.

BOOKKEEPING_FEATURE_TYPES `module-attribute` ¶

BOOKKEEPING_FEATURE_TYPES: dict[str, str] = {
    "response": "string",
    "finish_reason": "string",
    "num_tokens": "int64",
    "error": "string",
    "error_type": "string",
    "reasoning": "string",
    "valid_fields": "bool",
    "valid": "bool",
}

Arrow type of each bookkeeping column, keyed by its bare (unprefixed) name.

messages is not here: it holds the rendered prompt, of whatever type the input dataset gives it, and is typed from that dataset's features instead.

DEFAULT_MAX_CONCURRENT_BATCHES `module-attribute` ¶

DEFAULT_MAX_CONCURRENT_BATCHES = 4

Simultaneous batches per vLLM server unless a step says otherwise.

SelectionRecord `dataclass` ¶

SelectionRecord(
    max_num_samples: int | None,
    source_rows: int,
    selected_rows: int,
    components: dict[str, str] = dict(),
)

What a prepared dataset was built from.

prepare_data writes it to <output_dir>/<task_prefix>selection.json. The file stays in place after a run finishes, so a later run can tell whether the finished rows were annotated with the settings that it asks for.

Attributes:

Name	Type	Description
`max_num_samples`	`int \| None`	The requested sample cap, `None` for no cap. It is outside `components` because a run may grow under a higher cap.
`source_rows`	`int`	Number of rows in the source dataset.
`selected_rows`	`int`	Number of rows that the cap left in the selection.
`components`	`dict[str, str]`	One short string per setting that decides what the prepared dataset holds, keyed by the setting's name. A setting that the layer which owns it has not recorded yet is absent.

Examples:

record = SelectionRecord(
    max_num_samples=10,
    source_rows=40,
    selected_rows=10,
    components={"prompt_template": "abc", "shuffle_seed": "42"},
)
record.changed_components({"prompt_template": "abc"})
# []
record.changed_components({"prompt_template": "def"})
# ['prompt_template']
record.changed_components({"system_message": "def"})
# []

changed_components ¶

changed_components(components: dict[str, str]) -> list[str]

Name the requested settings that differ from the recorded ones.

Parameters:

Name	Type	Description	Default
`components`	`dict[str, str]`	The settings of the request, as `_preparation_components` builds them.	required

Returns:

Type	Description
`list[str]`	The names of the settings that the record holds under another
`list[str]`	value, sorted. A setting the record does not hold is left out:
`list[str]`	the layer that owns it writes it only once it ran, so the output
`list[str]`	schema is absent until `run_annotation` recorded it, and a
`list[str]`	step's input version until the pipeline recorded it.

View source on GitHub: src/llm_annotator/annotator.py lines 694–712

path `classmethod` ¶

path(output_dir: str | Path, task_prefix: str = '') -> Path

Build the path of the record file.

Parameters:

Name	Type	Description	Default
`output_dir`	`str \| Path`	The annotator's output directory.	required
`task_prefix`	`str`	The task prefix of the run.	`''`

Returns:

Type	Description
`Path`	Path of `<output_dir>/<task_prefix>selection.json`.

View source on GitHub: src/llm_annotator/annotator.py lines 714–725

read `classmethod` ¶

read(
    output_dir: str | Path, task_prefix: str = ""
) -> "SelectionRecord | None"

Read the record of an output directory.

Parameters:

Name	Type	Description	Default
`output_dir`	`str \| Path`	The annotator's output directory.	required
`task_prefix`	`str`	The task prefix of the run.	`''`

Returns:

Type	Description
`'SelectionRecord \| None'`	The record, or `None` when the directory has none, which is a
`'SelectionRecord \| None'`	run that never prepared data.

Raises:

Type	Description
`ValueError`	If the file holds no `components`, which is a record from a release that did not describe the settings of its run.

View source on GitHub: src/llm_annotator/annotator.py lines 727–765

write ¶

write(
    output_dir: str | Path, task_prefix: str = ""
) -> None

Write the record to an output directory.

Parameters:

Name	Type	Description	Default
`output_dir`	`str \| Path`	The annotator's output directory.	required
`task_prefix`	`str`	The task prefix of the run.	`''`

View source on GitHub: src/llm_annotator/annotator.py lines 767–776

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 handed to the client in one `batch_generate` call. It decides how much work is in flight at once, so it is what a provider's rate limit or a vLLM server's `max_num_seqs` has to cover.	`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 886–889

enter ¶

__enter__() -> 'Annotator'

Enter the context manager, returning the annotator instance.

View source on GitHub: src/llm_annotator/annotator.py lines 891–893

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 895–897

destroy ¶

destroy() -> None

Clean up all resources used by the underlying client.

View source on GitHub: src/llm_annotator/annotator.py lines 899–901

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,
    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,
    reuse_idx_column: bool = False,
    allow_selection_change: 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 applying the prompt template.	`None`
`idx_column`	`str`	Column name used as unique identifier. Must not exist in the input dataset.	`'idx'`
`task_prefix`	`str`	Prefix for the internal column names and for the artifacts of this task inside `output_dir`, so that several tasks can share one directory and one `hub_id`.	`''`
`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`
`reuse_idx_column`	`bool`	Whether an `idx_column` that already exists in the dataset is kept as the sample id. Use it for a dataset that an earlier run produced with `keep_idx_column=True`, so that a row keeps one id through several runs. A change of the source is then allowed, because the ids do not depend on the position of a row.	`False`
`allow_selection_change`	`bool`	Whether a request that the finished rows do not belong to is accepted. Only useful when those rows are deleted afterwards (`overwrite=True` in `run_annotation`).	`False`

Returns:

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

Raises:

Type	Description
`ValueError`	If progress files exist and the request no longer matches the settings that they were annotated with: an edited prompt template or system message, another `shuffle_seed`, a lower `max_num_samples`, or a source that changed in another way than appended rows without a shuffle.

View source on GitHub: src/llm_annotator/annotator.py lines 1593–1837

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,
    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
    | Literal["auto"] = "auto",
    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,
    retry_errors: bool | Sequence[str] = False,
) -> 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 discard the finished rows of this task and annotate every sample again. It removes `<output_dir>/<task_prefix>progress_backup/`, the Hub branch of the same name, `metadata/<task_prefix>annotation_metadata.json` and the final dataset in the root of `output_dir`. It keeps the prepared data (`<task_prefix>prepared_dataset/`, its Hub branch and `<task_prefix>selection.json`), so a crashed run resumes without preparing its data again; pass `force_data_preparation=True` to `prepare_data` to rebuild it. It also keeps the artifacts of every other `task_prefix` in the same directory.	`False`
`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`	JSON schema of the response, as a mapping or as JSON text. The client constrains the response to it, and every top-level property becomes a column. This is the only way to give a schema: `options.output_schema` is set from it and raises `ValueError` when the caller sets it.	`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 \| Literal['auto']`	Samples per JSONL progress file. `"auto"` is one percent of the rows with a floor of 1000, so at most 100 files are written and a resume stays cheap. A fixed number trades the samples lost at a crash against the cost of rescanning the files on every resume; 0 writes a single file of unlimited size.	`'auto'`
`task_prefix`	`str`	Prefix for the internal column names and for the artifacts of this task inside `output_dir`, so that several tasks can share one directory and one `hub_id`. The final dataset in the root of `output_dir` and on Hub `main` is shared by design: the task that finishes last replaces it, and with `keep_columns=True` it holds the columns of the tasks that ran before it.	`''`
`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. The rows of such batches are only written once a later batch succeeds, so a run that aborts leaves them to the resumed run. Set to 0 to disable.	`10`
`retry_errors`	`bool \| Sequence[str]`	Annotate rows again that finished with an error in an earlier run. By default such rows are final, so that an error caused by the sample (e.g. a prompt longer than the context) is not repeated on every resume. `True` redoes every errored row, a sequence redoes only those error types, e.g. `["ConnectError", "APITimeoutError"]`. The selected rows are removed from the progress files before the run.	`False`

Returns:

Type	Description
`Dataset`	Final concatenated annotation dataset.

Raises:

Type Description

ValueError

If no prepared data source can be resolved, if a top-level property of the schema has the name of a column that the annotator writes itself, if output_schema differs from the one that the finished rows were annotated with while overwrite is off, or if hub_id has a progress backup while the local progress directory is empty. Restore that backup with restore_progress_from_hub, or pass overwrite to discard it.

TooManyConsecutiveFailedBatchesError

If max_consecutive_failed_batches consecutive batches fail entirely.

View source on GitHub: src/llm_annotator/annotator.py lines 2280–2693

annotate_dataset ¶

annotate_dataset(
    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,
    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
    | Literal["auto"] = "auto",
    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,
    reuse_idx_column: bool = False,
    retry_errors: bool | Sequence[str] = False,
) -> 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. Every argument below is passed straight through, so the full description of each is in `prepare_data` (data selection and prompting) or in `run_annotation` (inference and output).	required
`prompt_template`	`str`	Prompt template with dataset fields.	required
`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`
`idx_column`	`str`	Column name used as the stable sample identifier.	`'idx'`
`task_prefix`	`str`	Prefix for this task's columns and artifacts.	`''`
`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 discard the finished rows of this task and annotate every sample again.	`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`.	`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 \| Literal['auto']`	Samples per JSONL progress file.	`'auto'`
`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.	`10`
`reuse_idx_column`	`bool`	Whether an `idx_column` that already exists in the dataset is kept as the sample id.	`False`
`retry_errors`	`bool \| Sequence[str]`	Annotate rows again that finished with an error in an earlier run.	`False`

Returns:

Type	Description
`Dataset`	The concatenated annotation dataset.

Raises:

Type	Description
`ValueError`	If finished rows exist that were annotated with other settings than the ones given now, and `overwrite` is off.
`TooManyConsecutiveFailedBatchesError`	If `max_consecutive_failed_batches` consecutive batches fail entirely.

View source on GitHub: src/llm_annotator/annotator.py lines 2695–2838

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
    | Literal["auto"] = "auto",
    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,
    retry_errors: bool | Sequence[str] = False,
) -> Dataset

Generate a new dataset from prompts.

Parameters:

Name	Type	Description	Default
`output_dir`	`str \| Path`	Directory where annotation output is written. `prompts` and `prompt_prefix` build the prompt dataset and its template; every other argument below is passed straight through, so the full description of each is in `prepare_data` (data selection and prompting) or in `run_annotation` (inference and output).	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 discard the finished rows of this task and annotate every sample again.	`False`
`options`	`ProviderRuntimeOptions \| None`	Runtime options passed to the client.	`None`
`gen_kwargs`	`dict[str, Any] \| None`	Extra request parameters merged over `options`.	`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 \| Literal['auto']`	Samples per JSONL progress file.	`'auto'`
`task_prefix`	`str`	Prefix for this task's columns and artifacts.	`''`
`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.	`10`
`retry_errors`	`bool \| Sequence[str]`	Annotate rows again that finished with an error in an earlier run.	`False`

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 2840–2953

push_progress_to_hub ¶

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

Upload the progress files of a run to the Hugging Face Hub.

Creates the dataset repository and its <task_prefix>progress_backup branch, and uploads the *.jsonl files of dir_path plus the selection record next to it, which is what llm_annotator.hub.restore_progress_from_hub reads back.

active_path names the progress file that the writer has open, so that a background upload does not read a file while it grows. That file is left out of the folder upload and its first active_bytes bytes are copied out and uploaded from the copy, which ends on a line boundary. Uploading the file itself would hash it and then read it again, and for a file that grew in between the stored object no longer matches the checksum of the commit.

Parameters:

Name	Type	Description	Default
`dir_path`	`Path \| str`	Directory that holds the `*.jsonl` progress files.	required
`hub_id`	`str`	Hugging Face dataset ID to upload into.	required
`task_prefix`	`str`	String prefix to use for branch naming.	`''`
`active_path`	`Path \| None`	The progress file that the writer has open, when the writer is running.	`None`
`active_bytes`	`int`	How many bytes of `active_path` were written when the upload was requested.	`0`

View source on GitHub: src/llm_annotator/annotator.py lines 3246–3330

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,
    max_workers: int | None = None,
    max_concurrent_batches_per_client: int = DEFAULT_MAX_CONCURRENT_BATCHES,
)

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 each client to a queue once per allowed concurrent batch;
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 concurrent batch slot, and any lower value is raised to the number of slots, since a smaller queue would leave servers idle. After initialisation the attribute always holds the resolved value. A config-driven run rejects a too-small value when the config loads instead of raising it here.	`None`
`max_concurrent_batches_per_client`	`int`	Maximum number of batches each server handles at once. This is independent of `batch_size`. Defaults to four for high throughput.	`DEFAULT_MAX_CONCURRENT_BATCHES`
`max_workers`	`int \| None`	Maximum worker threads used for batch annotation. It can exceed the initially available batch slots when additional servers are expected to join, allowing workers to wait for and immediately use those late-ready servers.	`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",
    )

is_shutting_down `property` ¶

is_shutting_down: bool

Whether the annotator has begun releasing its clients.

__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, `queue_size` is not positive, or `max_concurrent_batches_per_client` is not positive.
`TypeError`	If a client is not a vLLM server client.

View source on GitHub: src/llm_annotator/annotator.py lines 3461–3511

resolve_queue_size `staticmethod` ¶

resolve_queue_size(
    queue_size: int | None, num_slots: int
) -> int

Turn a requested queue size into the effective one for a pool.

This is the arithmetic behind the queue_size key, exposed so that a caller can report the effective value (as llm-annotate --describe-steps does) without building a pool first.

Parameters:

Name	Type	Description	Default
`queue_size`	`int \| None`	Requested number of batches in flight, or `None` to derive it from the number of slots.	required
`num_slots`	`int`	Concurrent batch slots in the pool, that is, servers times `max_concurrent_batches_per_client`.	required

Returns:

Type	Description
`int`	The number of batches to keep in flight: `num_slots` at the very
`int`	least, and `QUEUE_BATCHES_PER_SLOT` times that when nothing is
`int`	requested.

Raises:

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

Examples:

VLLMQueueAnnotator.resolve_queue_size(None, 8)
# 32
VLLMQueueAnnotator.resolve_queue_size(64, 8)
# 64
VLLMQueueAnnotator.resolve_queue_size(2, 8)
# 8

View source on GitHub: src/llm_annotator/annotator.py lines 3529–3567

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 3592–3608

add_client ¶

add_client(client: Client[Any]) -> None

Add a vLLM server that became ready after annotation started.

Parameters:

Name	Type	Description	Default
`client`	`Client[Any]`	Ready vLLM server client to make available to workers.	required

Raises:

Type	Description
`TypeError`	If `client` is not a vLLM server client.

View source on GitHub: src/llm_annotator/annotator.py lines 3610–3620

add_client_for_base_url ¶

add_client_for_base_url(
    base_url: str,
    client_factory: Callable[[str], Client[Any]],
) -> None

Construct and add a late-ready server under the pool lock.

View source on GitHub: src/llm_annotator/annotator.py lines 3622–3644

wait_for_shutdown ¶

wait_for_shutdown(timeout: float) -> bool

Block until the pool is shutting down or the timeout elapses.

View source on GitHub: src/llm_annotator/annotator.py lines 3676–3678

client_count ¶

client_count() -> int

Return the current number of pool members.

View source on GitHub: src/llm_annotator/annotator.py lines 3680–3683

client_base_urls ¶

client_base_urls() -> set[str]

Return the base URLs currently registered in the pool.

View source on GitHub: src/llm_annotator/annotator.py lines 3685–3692

set_max_concurrent_batches_per_client ¶

set_max_concurrent_batches_per_client(
    max_concurrent_batches_per_client: int,
) -> None

Change the concurrent-request limit for every server.

This may only be called between annotation runs because it rebuilds the available-client queue.

Parameters:

Name	Type	Description	Default
`max_concurrent_batches_per_client`	`int`	Maximum simultaneous batch requests sent to each server.	required

Raises:

Type	Description
`ValueError`	If the requested limit is not positive.
`RuntimeError`	If called while annotation is in progress.

View source on GitHub: src/llm_annotator/annotator.py lines 3694–3728

destroy ¶

destroy() -> None

Clean up the resources of every client in the pool.

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 3737–3762

is_retried_error ¶

is_retried_error(
    row: dict[str, Any],
    *,
    retry_errors: bool | Sequence[str],
    task_prefix: str = "",
) -> bool

Check whether retry_errors selects a finished row for another attempt.

Parameters:

Name	Type	Description	Default
`row`	`dict[str, Any]`	One row of a progress file.	required
`retry_errors`	`bool \| Sequence[str]`	`True` selects every errored row. A sequence selects the errored rows whose `error_type` is in it.	required
`task_prefix`	`str`	Prefix of the internal column names.	`''`

Returns:

Type	Description
`bool`	Whether the row has an error that `retry_errors` selects.

Examples:

row = {"error": "refused", "error_type": "ConnectError"}
is_retried_error(row, retry_errors=True)
# True
is_retried_error(row, retry_errors=["APITimeoutError"])
# False
is_retried_error({"error": None}, retry_errors=True)
# False

View source on GitHub: src/llm_annotator/annotator.py lines 463–495

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 779–808

Annotator¶

llm_annotator.annotator ¶

BOOKKEEPING_SUFFIXES module-attribute ¶

BOOKKEEPING_FEATURE_TYPES module-attribute ¶

DEFAULT_MAX_CONCURRENT_BATCHES module-attribute ¶

SelectionRecord dataclass ¶

changed_components ¶

path classmethod ¶

read classmethod ¶

write ¶

Annotator dataclass ¶

__post_init__ ¶

__enter__ ¶

__exit__ ¶

destroy ¶

prepare_data ¶

run_annotation ¶

annotate_dataset ¶

generate_dataset ¶

push_progress_to_hub ¶

VLLMQueueAnnotator dataclass ¶

is_shutting_down property ¶

__post_init__ ¶

resolve_queue_size staticmethod ¶

set_queue_size ¶

add_client ¶

add_client_for_base_url ¶

wait_for_shutdown ¶

client_count ¶

client_base_urls ¶

set_max_concurrent_batches_per_client ¶

destroy ¶

is_retried_error ¶

destroy_on_error ¶

BOOKKEEPING_SUFFIXES `module-attribute` ¶

BOOKKEEPING_FEATURE_TYPES `module-attribute` ¶

DEFAULT_MAX_CONCURRENT_BATCHES `module-attribute` ¶

SelectionRecord `dataclass` ¶

path `classmethod` ¶

read `classmethod` ¶

Annotator `dataclass` ¶

enter ¶

exit ¶

VLLMQueueAnnotator `dataclass` ¶

is_shutting_down `property` ¶

resolve_queue_size `staticmethod` ¶