Home

Awesome

<a href="https://explosion.ai"><img src="https://explosion.ai/assets/img/logo.svg" width="125" height="125" align="right" /></a>

spacy-huggingface-pipelines: Use pretrained transformer models for text and token classification

This package provides spaCy components to use pretrained Hugging Face Transformers pipelines for inference only.

PyPi GitHub

Features

🚀 Installation

Installing the package from pip will automatically install all dependencies, including PyTorch and spaCy.

pip install -U pip setuptools wheel
pip install spacy-huggingface-pipelines

For GPU installation, follow the spaCy installation quickstart with GPU, e.g.

pip install -U spacy[cuda12x]

If you are having trouble installing PyTorch, follow the instructions on the official website for your specific operating system and requirements.

📖 Documentation

This module provides spaCy wrappers for the inference-only transformers TokenClassificationPipeline and TextClassificationPipeline pipelines.

The models are downloaded on initialization from the Hugging Face Hub if they're not already in your local cache, or alternatively they can be loaded from a local path.

Note that the transformer model data is not saved with the pipeline when you call nlp.to_disk, so if you are loading pipelines in an environment with limited internet access, make sure the model is available in your transformers cache directory and enable offline mode if needed.

Token classification

Config settings for hf_token_pipe:

[components.hf_token_pipe]
factory = "hf_token_pipe"
model = "dslim/bert-base-NER"     # Model name or path
revision = "main"                 # Model revision
aggregation_strategy = "average"  # "simple", "first", "average", "max"
stride = 16                       # If stride >= 0, process long texts in
                                  # overlapping windows of the model max
                                  # length. The value is the length of the
                                  # window overlap in transformer tokenizer
                                  # tokens, NOT the length of the stride.
kwargs = {}                       # Any additional arguments for
                                  # TokenClassificationPipeline
alignment_mode = "strict"         # "strict", "contract", "expand"
annotate = "ents"                 # "ents", "pos", "spans", "tag"
annotate_spans_key = null         # Doc.spans key for annotate = "spans"
scorer = null                     # Optional scorer

TokenClassificationPipeline settings

spaCy settings

Examples

  1. Save named entity annotation as Doc.ents:
import spacy
nlp = spacy.blank("en")
nlp.add_pipe("hf_token_pipe", config={"model": "dslim/bert-base-NER"})
doc = nlp("My name is Sarah and I live in London")
print(doc.ents)
# (Sarah, London)
  1. Save named entity annotation as Doc.spans[spans_key] and scores as Doc.spans[spans_key].attrs["scores"]:
import spacy
nlp = spacy.blank("en")
nlp.add_pipe(
    "hf_token_pipe",
    config={
        "model": "dslim/bert-base-NER",
        "annotate": "spans",
        "annotate_spans_key": "bert-base-ner",
    },
)
doc = nlp("My name is Sarah and I live in London")
print(doc.spans["bert-base-ner"])
# [Sarah, London]
print(doc.spans["bert-base-ner"].attrs["scores"])
# [0.99854773, 0.9996215]
  1. Save fine-grained tags as Token.tag:
import spacy
nlp = spacy.blank("en")
nlp.add_pipe(
    "hf_token_pipe",
    config={
        "model": "QCRI/bert-base-multilingual-cased-pos-english",
        "annotate": "tag",
    },
)
doc = nlp("My name is Sarah and I live in London")
print([t.tag_ for t in doc])
# ['PRP$', 'NN', 'VBZ', 'NNP', 'CC', 'PRP', 'VBP', 'IN', 'NNP']
  1. Save coarse-grained tags as Token.pos:
import spacy
nlp = spacy.blank("en")
nlp.add_pipe(
    "hf_token_pipe",
    config={"model": "vblagoje/bert-english-uncased-finetuned-pos", "annotate": "pos"},
)
doc = nlp("My name is Sarah and I live in London")
print([t.pos_ for t in doc])
# ['PRON', 'NOUN', 'AUX', 'PROPN', 'CCONJ', 'PRON', 'VERB', 'ADP', 'PROPN']

Text classification

Config settings for hf_text_pipe:

[components.hf_text_pipe]
factory = "hf_text_pipe"
model = "distilbert-base-uncased-finetuned-sst-2-english"  # Model name or path
revision = "main"                 # Model revision
kwargs = {}                       # Any additional arguments for
                                  # TextClassificationPipeline
scorer = null                     # Optional scorer

The input texts are truncated according to the transformers model max length.

TextClassificationPipeline settings

Example

import spacy

nlp = spacy.blank("en")
nlp.add_pipe(
    "hf_text_pipe",
    config={"model": "distilbert-base-uncased-finetuned-sst-2-english"},
)
doc = nlp("This is great!")
print(doc.cats)
# {'POSITIVE': 0.9998694658279419, 'NEGATIVE': 0.00013048505934420973}

Batching and GPU

Both token and text classification support batching with nlp.pipe:

for doc in nlp.pipe(texts, batch_size=256):
    do_something(doc)

If the component runs into an error processing a batch (e.g. on an empty text), nlp.pipe will back off to processing each text individually. If it runs into an error on an individual text, a warning is shown and the doc is returned without additional annotation.

Switch to GPU:

import spacy
spacy.require_gpu()

for doc in nlp.pipe(texts):
    do_something(doc)

Bug reports and issues

Please report bugs in the spaCy issue tracker or open a new thread on the discussion board for other issues.