nlu.annotation.joint_bert.joint_bert

Joint BERT model for intent classification and slot annotation.

Module Contents

Classes

JointBERT

class nlu.annotation.joint_bert.joint_bert.JointBERT(intent_label_count: int, slot_label_count: int)

Bases: torch.nn.Module

forward(input_ids: torch.Tensor, attention_mask: torch.Tensor) → Tuple[torch.Tensor, torch.Tensor]

Forward pass of the model.

Parameters:

• input_ids – The input token IDs.

• attention_mask – The attention mask.

Returns:

Tuple of intent and slot logits.

predict(input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None) → Tuple[int, List[int]]

Predicts the intent and slot annotations for the given input.

Parameters:

• input_ids – The input token IDs.

• attention_mask (optional) – The attention mask. Defaults to None.

Returns:

A tuple of the predicted intent and slot annotations.

classmethod from_pretrained(path: str) → JointBERT

Loads the model and tokenizer from the specified directory.

Parameters:

path – The path to the directory containing the model and tokenizer.

Returns:

The loaded model.