logitorch.models.bertnot

Classes

BERTNOT

BERTNOT model for fine-tuning BERT for various tasks.

Module Contents

class logitorch.models.bertnot.BERTNOT(pretrained_bert_model: str, num_labels: int = 2)

Bases: torch.nn.Module

BERTNOT model for fine-tuning BERT for various tasks.

Args:

pretrained_bert_model (str): Path or identifier of the pre-trained BERT model. num_labels (int, optional): Number of labels for the classification task. Defaults to 2.

forward(x, y=None, task='mlm', loss='cross_entropy')

Forward pass of the BERTNOT model.

Args:

x (dict): Input dictionary containing the input tensors. y (torch.Tensor, optional): Target tensor. Defaults to None. task (str, optional): Task type. Defaults to “mlm”. loss (str, optional): Loss type. Defaults to “cross_entropy”.

Returns:

tuple: Tuple containing the loss and logits if y is not None, otherwise returns logits.

predict(context: str, hypothesis: str = None, task='mlm', device='cpu')

Perform prediction using the BERTNOT model.

Args:

context (str): Input context string. hypothesis (str, optional): Input hypothesis string. Defaults to None. task (str, optional): Task type. Defaults to “mlm”. device (str, optional): Device to run the model on. Defaults to “cpu”.

Returns:

str or int: Predicted token or label.

cross_entopy_loss

dropout

kl_loss

log_softmax

losses = ['cross_entropy', 'unlikelihood', 'kl']

model

num_labels = 2

original_bert

original_bert_softmax

sequence_classifier

tasks = ['mlm', 'te']

tokenizer

unlikelihood_loss