nessie.models.text.transformer_text_classifier

Module Contents

Classes

TextClassificationDataset
TextClassificationEvalDataset
TransformerTextClassifier	Helper class that provides a standard way to create an ABC using

Functions

compute_metrics(pred)

class nessie.models.text.transformer_text_classifier.TextClassificationDataset(tokenized_texts: Dict, encoded_labels: List[int])

Bases: torch.utils.data.Dataset

__getitem__(self, idx)

__len__(self)

class nessie.models.text.transformer_text_classifier.TextClassificationEvalDataset(tokenized_texts: Dict)

Bases: torch.utils.data.Dataset

__getitem__(self, idx)

__len__(self)

class nessie.models.text.transformer_text_classifier.TransformerTextClassifier(verbose: bool = True, max_epochs: int = 24, batch_size: int = 16, model_name: str = BERT_BASE)

Bases: nessie.models.model.TextClassifier, nessie.models.model.Callbackable

Helper class that provides a standard way to create an ABC using inheritance.

add_callback(self, name: str, callback: transformers.TrainerCallback)

fit(self, X: nessie.types.StringArray, y: nessie.types.StringArray)

has_dropout(self) → bool

label_encoder(self) → sklearn.preprocessing.LabelEncoder

Returns a label encoder that can be used to map labels to ints and vice versa

predict(self, X: nessie.types.StringArray) → numpy.typing.NDArray[str]

predict_proba(self, X: nessie.types.StringArray) → numpy.typing.NDArray[float]

Returns a distribution over labels for each instance.

Parameters

X – The texts to predict on

Returns

A (num_instances, num_labels) numpy array

score(self, X: nessie.types.StringArray) → numpy.typing.NDArray[float]

Returns the best score for each item

use_dropout(self, is_activated: bool)

nessie.models.text.transformer_text_classifier.compute_metrics(pred)