|
| 1 | +--- |
| 2 | +title: "TransformersNamedEntityExtractor" |
| 3 | +id: transformersnamedentityextractor |
| 4 | +slug: "/transformersnamedentityextractor" |
| 5 | +description: "This component extracts predefined entities out of a piece of text and writes them into documents’ meta field." |
| 6 | +--- |
| 7 | + |
| 8 | +# TransformersNamedEntityExtractor |
| 9 | + |
| 10 | +This component extracts predefined entities out of a piece of text and writes them into documents’ meta field. |
| 11 | + |
| 12 | +<div className="key-value-table"> |
| 13 | + |
| 14 | +| | | |
| 15 | +| --- | --- | |
| 16 | +| **Most common position in a pipeline** | After the [PreProcessor](../preprocessors.mdx) in an indexing pipeline or after a [Retriever](../retrievers.mdx) in a query pipeline | |
| 17 | +| **Mandatory init variables** | `model`: Name or path of the model to use | |
| 18 | +| **Mandatory run variables** | `documents`: A list of documents | |
| 19 | +| **Output variables** | `documents`: A list of documents | |
| 20 | +| **API reference** | [Transformers](/reference/integrations-transformers) | |
| 21 | +| **GitHub link** | https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/transformers | |
| 22 | +| **Package name** | `transformers-haystack` | |
| 23 | + |
| 24 | +</div> |
| 25 | + |
| 26 | +## Overview |
| 27 | + |
| 28 | +`TransformersNamedEntityExtractor` looks for entities, which are spans in the text. The extractor automatically recognizes and groups them depending on their class, such as people's names, organizations, locations, and other types. The exact classes are determined by the model that you initialize the component with. |
| 29 | + |
| 30 | +`TransformersNamedEntityExtractor` takes a list of documents as input and returns a list of the same documents with their `meta` data enriched with `NamedEntityAnnotations`. A `NamedEntityAnnotation` consists of the type of the entity, the start and end of the span, and a score calculated by the model, for example: `NamedEntityAnnotation(entity='PER', start=11, end=16, score=0.9)`. |
| 31 | + |
| 32 | +When the `TransformersNamedEntityExtractor` is initialized, you need to set a `model`. Optionally, you can set `pipeline_kwargs`, which are then passed on to the Hugging Face pipeline. You can additionally set the `device` that is used to run the component. |
| 33 | + |
| 34 | +Authentication with a Hugging Face API token is only required to access private or gated models. You can pass the token at initialization with `token`, or set the `HF_API_TOKEN` or `HF_TOKEN` environment variable. |
| 35 | + |
| 36 | +## Usage |
| 37 | + |
| 38 | +Install the `transformers-haystack` package to use the `TransformersNamedEntityExtractor`: |
| 39 | + |
| 40 | +```shell |
| 41 | +pip install transformers-haystack |
| 42 | +``` |
| 43 | + |
| 44 | +The component works with any Hugging Face model that supports token classification or NER. |
| 45 | + |
| 46 | +`TransformersNamedEntityExtractor` accepts a list of `Documents` as its input. The extractor annotates the raw text in the documents and stores the annotations in the document's `meta` dictionary under the `named_entities` key. |
| 47 | + |
| 48 | +```python |
| 49 | +from haystack.dataclasses import Document |
| 50 | +from haystack_integrations.components.extractors.transformers import ( |
| 51 | + TransformersNamedEntityExtractor, |
| 52 | +) |
| 53 | + |
| 54 | +extractor = TransformersNamedEntityExtractor(model="dslim/bert-base-NER") |
| 55 | + |
| 56 | +documents = [ |
| 57 | + Document(content="My name is Clara and I live in Berkeley, California."), |
| 58 | + Document(content="I'm Merlin, the happy pig!"), |
| 59 | + Document(content="New York State is home to the Empire State Building."), |
| 60 | +] |
| 61 | + |
| 62 | +extractor.run(documents) |
| 63 | +print(documents) |
| 64 | +``` |
| 65 | + |
| 66 | +Here is the example result: |
| 67 | + |
| 68 | +```python |
| 69 | +[Document(id=aec840d1b6c85609f4f16c3e222a5a25fd8c4c53bd981a40c1268ab9c72cee10, content: 'My name is Clara and I live in Berkeley, California.', meta: {'named_entities': [NamedEntityAnnotation(entity='PER', start=11, end=16, score=0.99641764), NamedEntityAnnotation(entity='LOC', start=31, end=39, score=0.996198), NamedEntityAnnotation(entity='LOC', start=41, end=51, score=0.9990196)]}), |
| 70 | +Document(id=98f1dc5d0ccd9d9950cd191d1076db0f7af40c401dd7608f11c90cb3fc38c0c2, content: 'I'm Merlin, the happy pig!', meta: {'named_entities': [NamedEntityAnnotation(entity='PER', start=4, end=10, score=0.99054915)]}), |
| 71 | +Document(id=44948ea0eec018b33aceaaedde4616eb9e93ce075e0090ec1613fc145f84b4a9, content: 'New York State is home to the Empire State Building.', meta: {'named_entities': [NamedEntityAnnotation(entity='LOC', start=0, end=14, score=0.9989541), NamedEntityAnnotation(entity='LOC', start=30, end=51, score=0.95746297)]})] |
| 72 | +``` |
| 73 | + |
| 74 | +### Get stored annotations |
| 75 | + |
| 76 | +This component includes the `get_stored_annotations` helper class method that allows you to retrieve the annotations stored in a `Document` transparently: |
| 77 | + |
| 78 | +```python |
| 79 | +from haystack.dataclasses import Document |
| 80 | +from haystack_integrations.components.extractors.transformers import ( |
| 81 | + TransformersNamedEntityExtractor, |
| 82 | +) |
| 83 | + |
| 84 | +extractor = TransformersNamedEntityExtractor(model="dslim/bert-base-NER") |
| 85 | + |
| 86 | +documents = [ |
| 87 | + Document(content="My name is Clara and I live in Berkeley, California."), |
| 88 | + Document(content="I'm Merlin, the happy pig!"), |
| 89 | + Document(content="New York State is home to the Empire State Building."), |
| 90 | +] |
| 91 | + |
| 92 | +extractor.run(documents) |
| 93 | + |
| 94 | +annotations = [ |
| 95 | + TransformersNamedEntityExtractor.get_stored_annotations(doc) for doc in documents |
| 96 | +] |
| 97 | +print(annotations) |
| 98 | + |
| 99 | +# If a Document doesn't contain any annotations, this returns None. |
| 100 | +new_doc = Document(content="In one of many possible worlds...") |
| 101 | +assert TransformersNamedEntityExtractor.get_stored_annotations(new_doc) is None |
| 102 | +``` |
0 commit comments