spacybert: Bert inference for spaCy

spaCy v2.0 extension and pipeline component for loading BERT sentence / document embedding meta data to Doc, Span and Token objects. The Bert backend itself is supported by the Hugging Face transformers library.

Installation

spacybert requires spacy v2.0.0 or higher.

Usage

Getting BERT embeddings for single language dataset

import spacy
from spacybert import BertInference
nlp = spacy.load('en')

Then either use BertInference as part of a pipeline,

bert = BertInference(
    from_pretrained='path/to/pretrained_bert_weights_dir',
    set_extension=False)
nlp.add_pipe(bert, last=True)

Or not...

bert = BertInference(
    from_pretrained='path/to/pretrained_bert_weights_dir',
    set_extension=True)

The difference is that when set_extension=True, bert_repr is set as a property extension for the Doc, Span and Token spacy objects. If set_extension=False, the bert_repr is set as an attribute extension with a default value (=None). The attribute computes the correct value when doc._.bert_repr is called.

Get the Bert representation / embedding.

doc = nlp("This is a test")
print(doc._.bert_repr)  # <-- torch.Tensor

Getting BERT embeddings for multiple languages dataset.

import spacy
from spacy_langdetect import LanguageDetector
from spacybert import MultiLangBertInference

nlp = spacy.load('en')
nlp.add_pipe(LanguageDetector(), name='language_detector', last=True)
bert = MultiLangBertInference(
    from_pretrained={
        'en': 'path/to/en_pretrained_bert_weights_dir',
        'nl': 'path/to/nl_pretrained_bert_weights_dir'
    },
    set_extension=False)
nlp.add_pipe(bert, after='language_detector')

texts = [
    "This is a test",  # English
    "Dit is een test"  # Dutch
]
for doc in nlp.pipe(texts):
    print(doc._.bert_repr)  # <-- torch.Tensor

When language_detector detects languages other than the ones for which pre-trained weights is specified, by default doc._.bert_repr = None.

Available attributes

The extension sets attributes on the Doc, Span and Token. You can change the attribute name on initializing the extension.

Settings

On initialization of BertInference, you can define the following:

On initialization of MultiLangBertInference, you can define the following:

Pooling strategies

Roadmap

This extension is still experimental. Possible future updates include:

• Getting document representation from other state-of-the-art NLP models other than Google's BERT.
• Method for computing similarity between Doc, Span and Token objects using the bert_repr tensor.
• Getting representation from multiple / other layers in the models.