tfm.nlp.models.DualEncoder

A dual encoder model based on a transformer-based encoder.

tfm.nlp.models.DualEncoder(
    network: tf_keras.Model,
    max_seq_length: int = 32,
    normalize: bool = True,
    logit_scale: float = 1.0,
    logit_margin: float = 0.0,
    output: str = 'logits',
    **kwargs
) -> None

This is an implementation of the dual encoder network structure based on the transfomer stack, as described in "Language-agnostic BERT Sentence Embedding"

The DualEncoder allows a user to pass in a transformer stack, and build a dual encoder model based on the transformer stack.

network A transformer network which should output an encoding output.
max_seq_length The maximum allowed sequence length for transformer.
normalize If set to True, normalize the encoding produced by transfomer.
logit_scale The scaling factor of dot products when doing training.
logit_margin The margin between positive and negative when doing training.
output The output style for this network. Can be either logits or predictions. If set to predictions, it will output the embedding producted by transformer network.

checkpoint_items Returns a dictionary of items to be additionally checkpointed.

Methods

call

call(
    inputs, training=None, mask=None
)

Calls the model on new inputs and returns the outputs as tensors.

In this case call() just reapplies all ops in the graph to the new inputs (e.g. build a new computational graph from the provided inputs).

Args
inputs Input tensor, or dict/list/tuple of input tensors.
training Boolean or boolean scalar tensor, indicating whether to run the Network in training mode or inference mode.
mask A mask or list of masks. A mask can be either a boolean tensor or None (no mask). For more details, check the guide here.

Returns
A tensor if there is a single output, or a list of tensors if there are more than one outputs.