classifiers.prompt_task_complexity

Module Contents

Classes

CustomHFDeberta
MeanPooling
MulticlassHead
PromptTaskComplexityClassifier
PromptTaskComplexityConfig
PromptTaskComplexityModel

Data

PROMPT_TASK_COMPLEXITY_IDENTIFIER

API

class classifiers.prompt_task_complexity.CustomHFDeberta(config: dataclasses.dataclass)

Bases: torch.nn.Module, huggingface_hub.PyTorchModelHubMixin

Initialization

compute_results(

preds: torch.Tensor,

target: str,

decimal: int = 4,

) → tuple[list[str], list[str], list[float]]

forward(batch: dict[str, torch.Tensor]) → dict[str, torch.Tensor]

process_logits(logits: list[torch.Tensor]) → dict[str, torch.Tensor]

set_autocast(autocast: bool) → None

class classifiers.prompt_task_complexity.MeanPooling

Bases: torch.nn.Module

Initialization

forward(

last_hidden_state: torch.Tensor,

attention_mask: torch.Tensor,

) → torch.Tensor

class classifiers.prompt_task_complexity.MulticlassHead(input_size: int, num_classes: int)

Bases: torch.nn.Module

Initialization

forward(x: torch.Tensor) → torch.Tensor

classifiers.prompt_task_complexity.PROMPT_TASK_COMPLEXITY_IDENTIFIER

‘nvidia/prompt-task-and-complexity-classifier’

class classifiers.prompt_task_complexity.PromptTaskComplexityClassifier(

batch_size: int = 256,

text_field: str = 'text',

max_chars: int = 2000,

device_type: str = 'cuda',

autocast: bool = True,

max_mem_gb: int | None = None,

)

Bases: nemo_curator.classifiers.base.DistributedDataClassifier

Initialization

get_labels() → list[str]

class classifiers.prompt_task_complexity.PromptTaskComplexityConfig

base_model: str

‘microsoft/DeBERTa-v3-base’

max_len: int

512

model_output_type: dict

‘field(…)’

class classifiers.prompt_task_complexity.PromptTaskComplexityModel(

config: classifiers.prompt_task_complexity.PromptTaskComplexityConfig,

autocast: bool,

max_mem_gb: int | None,

)

Bases: crossfit.backend.torch.hf.model.HFModel

Initialization

load_config() → transformers.AutoConfig

load_model(

device: str = 'cuda',

) → classifiers.prompt_task_complexity.CustomHFDeberta

load_tokenizer() → transformers.AutoTokenizer