speechbrain.lobes.models.huggingface_transformers.textencoder module

This lobe enables the integration of generic huggingface pretrained text encoders (e.g. BERT).

Transformer from HuggingFace needs to be installed: https://huggingface.co/transformers/installation.html

Authors

  • Sylvain de Langen 2024

Summary

Classes:

TextEncoder

This lobe enables the integration of a generic HuggingFace text encoder (e.g. BERT).

Reference

class speechbrain.lobes.models.huggingface_transformers.textencoder.TextEncoder(source, save_path, freeze=True, num_layers: int | None = None, **kwargs)[source]

Bases: HFTransformersInterface

This lobe enables the integration of a generic HuggingFace text encoder (e.g. BERT). Requires the AutoModel found from the source to have a last_hidden_state key in the output dict.

Parameters:

  • source (str) – HuggingFace hub name: e.g β€œgoogle-bert/bert-base”

  • save_path (str) – Path (dir) of the downloaded model.

  • freeze (bool (default: True)) – If True, the model is frozen. If False, the model will be trained alongside with the rest of the pipeline.

  • num_layers (int, optional) – When specified, and assuming the passed LM can be truncated that way, the encoder for the passed model will be truncated to the specified layer (mutating it). This means that the embeddings will be those of the Nth layer rather than the last layer. The last layer is not necessarily the best for certain tasks.

  • **kwargs – Extra keyword arguments passed to the from_pretrained function.

Example

>>> inputs = ["La vie est belle"]
>>> model_hub = "google-bert/bert-base-multilingual-cased"
>>> save_path = "savedir"
>>> model = TextEncoder(model_hub, save_path)
>>> outputs = model(inputs)
truncate(keep_layers: int)[source]

Truncates the encoder to a specific layer so that output embeddings are the hidden state of the n-th layer.

Parameters:

keep_layers (int) – Number of layers to keep, e.g. 4 would keep layers [0, 1, 2, 3].

forward(input_texts, return_tokens: bool = False)[source]

This method implements a forward of the encoder model, which generates batches of embeddings embeddings from input text.

Parameters:

  • input_texts (list of str) – The list of texts (required).

  • return_tokens (bool) – Whether to also return the tokens.

Returns:

  • (any, torch.Tensor) if return_tokens == True – Respectively: - Tokenized sentence in the form of a padded batch tensor. In the HF

    format, as returned by the tokenizer.

    • Output embeddings of the model (i.e. the last hidden state)

  • torch.Tensor if return_tokens == False – Output embeddings of the model (i.e. the last hidden state)