nlu.annotation.joint_bert.joint_bert_train

Training script for the JointBERT model.

Module Contents

Classes

JointBERTTrain

Functions

parse_arguments()

Parses the command line arguments.

Attributes

Batch

args

nlu.annotation.joint_bert.joint_bert_train.Batch
class nlu.annotation.joint_bert.joint_bert_train.JointBERTTrain(intent_label_count: int, slot_label_count: int, **kwargs)

Bases: moviebot.nlu.annotation.joint_bert.JointBERT, pytorch_lightning.LightningModule

training_step(batch: Batch, batch_idx: int) torch.Tensor

Training step for the JointBERT model.

Parameters:
  • batch – A batch of data.

  • batch_idx – Index of the batch.

Returns:

The loss for the batch.

validation_step(batch: Batch, batch_idx: int) torch.Tensor

Validation step for the JointBERT model.

Parameters:
  • batch – A batch of data.

  • batch_idx – Index of the batch.

Returns:

The loss for the batch.

configure_optimizers() Tuple[List, List]

Configures the optimizer and scheduler for training.

save_pretrained(path: str) None

Saves the model to the specified directory.

nlu.annotation.joint_bert.joint_bert_train.parse_arguments()

Parses the command line arguments.

nlu.annotation.joint_bert.joint_bert_train.args