task_models

class datadreamer.task_models.TaskModel(cache_folder_path=None)

Bases: _Cachable

unload_model()

class datadreamer.task_models.HFClassificationTaskModel(model_name, revision=None, trust_remote_code=False, device=None, device_map=None, dtype=None, adapter_name=None, adapter_kwargs=None, cache_folder_path=None, **kwargs)

Bases: TaskModel

property model: PreTrainedModel

property tokenizer: PreTrainedTokenizer

property model_max_length: int

run(texts, truncate=False, batch_size=10, batch_scheduler_buffer_size=None, adaptive_batch_size=True, progress_interval=60, force=False, cache_only=False, verbose=None, log_level=None, total_num_texts=None, return_generator=False, **kwargs)

Return type:

Union[Generator[dict[str, float], None, None], list[dict[str, float]]]

class datadreamer.task_models.ParallelTaskModel(*task_models)

Bases: _ParallelCachable, TaskModel

run(texts, *args, **kwargs)

Return type:

Union[Generator[str | list[str], None, None], list[str | list[str]]]