Use kb_ids in doc as training data for Entity Linker

[alfredomg]

Hello,

I'm trying to train a custom EL based on a trained custom NER using the new Spacy 3 config file system. The docs in my training DocBin file have been manually annotated with entities and their kb_id_'s as per my KB. However, when attempting to train EL I get the following warning:

[...] entity_linker.py:252: UserWarning: [W093] Could not find any data to train the Entity Linker on. Is your input data correctly formatted?

And there seems to be no training going on at all as EL scores stay at 0.

The config file I'm using is somewhat similar to the example one at https://github.com/explosion/projects/blob/v3/tutorials/nel_emerson/configs/nel.cfg

I think the problem boils down to the get_candidates function in the [components.entity_linker] option. I tried using the default get_candidates = {"@misc":"spacy.CandidateGenerator.v1"}, removing it altogether, as well as specifying my own custom registered function. They all resulted in the same behaviour.

Is there a way to force the EL trainer to just use the kb_ids that are annotated in the DocBin docs? Notice that I don't really need to use any "candidate generators" as I already know the correct kb_id for each entity in my training set.

FYI, this is the custom registered function I used:

@registry.misc("EntityIDAnnotations.v1")
def create_candidates():
    def get_entity_id_annotations(kb, span):
        return list(filter(lambda cand: cand.entity_ == span.kb_id_, kb.get_alias_candidates(span.text)))
    return get_entity_id_annotations

I don't like this implementation, as it should instead generate a list of Candidates directly. But as I said, I shouldn't even need to do that as I already know the correct kb_id, i.e. there will be one and only one "candidate" kb_id for every entity in my training set.

And this is the config file I used (using the above registered function):

[paths]
train = null
dev = null
kb = null
init_tok2vec = null
raw_text = null
base_nlp = null
vectors = "${paths.base_nlp}"

[system]
gpu_allocator = null
seed = 0

[nlp]
lang = "en"
pipeline = ["sentencizer","ner","entity_linker"]
disabled = []
before_creation = null
after_creation = null
after_pipeline_creation = null
tokenizer = {"@tokenizers":"spacy.Tokenizer.v1"}

[components]

[components.sentencizer]
factory = "sentencizer"
punct_chars = null

[components.ner]
source = "${paths.base_nlp}"
component = "ner"

[components.entity_linker]
factory = "entity_linker"
entity_vector_length = 64
get_candidates = {"@misc":"EntityIDAnnotations.v1"}
incl_context = true
incl_prior = true
labels_discard = []
n_sents = 0

[components.entity_linker.model]
@architectures = "spacy.EntityLinker.v1"
nO = null

[components.entity_linker.model.tok2vec]
@architectures = "spacy.HashEmbedCNN.v1"
pretrained_vectors = null
width = 96
depth = 2
embed_size = 2000
window_size = 1
maxout_pieces = 3
subword_features = true


[corpora]

[corpora.dev]
@readers = "spacy.Corpus.v1"
path = ${paths.dev}
gold_preproc = true

[corpora.train]
@readers = "spacy.Corpus.v1"
path = ${paths.train}
gold_preproc = true

[training]
dev_corpus = "corpora.dev"
train_corpus = "corpora.train"
seed = ${system.seed}
gpu_allocator = ${system.gpu_allocator}
dropout = 0.2
accumulate_gradient = 1
patience = 10000
max_epochs = 0
max_steps = 40000
eval_frequency = 200
frozen_components = ["sentencizer", "ner"]
before_to_disk = null

[training.batcher]
@batchers = "spacy.batch_by_words.v1"
discard_oversize = false
tolerance = 0.2
get_length = null

[training.batcher.size]
@schedules = "compounding.v1"
start = 100
stop = 1000
compound = 1.001
t = 0.0

[training.logger]
@loggers = "spacy.ConsoleLogger.v1"
progress_bar = false

[training.optimizer]
@optimizers = "Adam.v1"
beta1 = 0.9
beta2 = 0.999
L2_is_weight_decay = true
L2 = 0.01
grad_clip = 1.0
use_averages = false
eps = 0.00000001
learn_rate = 0.001

[training.score_weights]
ents_f = 0.0
ents_p = 0.0
ents_r = 1.0
ents_per_type = null
nel_micro_f = 0.5
nel_micro_r = null
nel_micro_p = null

[pretraining]
max_epochs = 1000
dropout = 0.2
n_save_every = null
component = "tok2vec"
layer = ""
corpus = "corpora.pretrain"

[pretraining.batcher]
@batchers = "spacy.batch_by_words.v1"
size = 3000
discard_oversize = false
tolerance = 0.2
get_length = null

[pretraining.objective]
@architectures = "spacy.PretrainCharacters.v1"
maxout_pieces = 3
hidden_size = 300
n_characters = 4

[pretraining.optimizer]
@optimizers = "Adam.v1"
beta1 = 0.9
beta2 = 0.999
L2_is_weight_decay = true
L2 = 0.01
grad_clip = 1.0
use_averages = true
eps = 0.00000001
learn_rate = 0.001

[initialize]
vectors = "${paths.base_nlp}"
init_tok2vec = ${paths.init_tok2vec}
vocab_data = null
lookups = null
before_init = null
after_init = null

[initialize.components]

[initialize.components.entity_linker]

[initialize.components.entity_linker.kb_loader]
@misc = "spacy.KBFromFile.v1"
kb_path = ${paths.kb}

[initialize.tokenizer]

And this is the training command:

$ python -m spacy train myconfig.cfg --paths.train ./trn_docs.spacy --paths.dev ./dev_docs.spacy --paths.kb ./my_kb.kb --paths.base_nlp ./my_ner/model-last --output ./my_el --code el_annotations.py

Thanks in advance for your help!

Alfredo

[polm]

Sorry you're having trouble with this. It sounds like the annotations may not be getting picked up for some reason. The Entity Linker code is in pure Python, so you should be able to debug it in place to figure out what's going on. If you look at this loop in the source you should be able to figure out what's going on - in particular, you should be able to see if kb_ids is empty for some reason. (You can just print some values as you go over examples in the loop.) Maybe you aren't setting the attributes the way you think you are?

Is there a way to force the EL trainer to just use the kb_ids that are annotated in the DocBin docs? Notice that I don't really need to use any "candidate generators" as I already know the correct kb_id for each entity in my training set.

Your model needs to have a way to generate candidates at inference time, right? Even if you just have a fixed list of candidates you want it to pick from, you need a function to provide that list. If you just use the annotated kb_ids your model will learn to just pick the ID you gave it, which is not meaningful behavior.

I'm a little surprised that changing the candidate generator has no effect on training.

[alfredomg]

Thanks for the quick reply. I'll try debugging and will report what I find.

[alfredomg]

Thanks again for the reply. I have tried what I describe below in the meantime. I'm trying this with toy data so that I can share the whole process here. I'm not getting the warning with this toy data, but the NEL_MICRO_F scores stay at 0 (so no learning is taking place).

My actions were prompted by your question:

Maybe you aren't setting the attributes the way you think you are?

That might be the case. Maybe I'm doing that wrong. So I decided to show you how I set the attributes (annotate the docs), train my NER, build my KB and attempt to train EL with toy data:

This is how I annotate the docs:

Assuming my training data has the format dataset = [[text, [[start_char, end_char, ner_label, kb_id_], ...]]], ...], e.g. toy data:

dataset = [
  ["The price of gas in Europe could increase due to pressure by Gazprom and the Russian state.", [[13, 16, "ENERGY", "C01"]]],
  ["Lead was used to increase the octane range in gas before the invention of unleaded gasoline.", [[0, 4, "MINERAL", "C02"], [46, 49, "ENERGY", "C03"], [83, 91, "ENERGY", "C03"]]],
  ["ACME's range of industrial ovens run on gas.", [[40, 43, "ENERGY", "C01"]]],
  ["The prices of maize and beans are linked to the price of gasoline.", [[14, 19, "AGRI", "C04"], [24, 29, "AGRI", "C05"], [57, 65, "ENERGY", "C03"]]]
]

I convert it to DocBin using this code:

nlp = spacy.load("en_core_web_lg")
docbin = DocBin(store_user_data=True)
docs = nlp.pipe((text for text, _ in dataset))

for doc, (_, anns) in zip(docs, dataset):
    ents = [
        doc.char_span(
            start_char_ix, end_char_ix, label=label, kb_id=kb_id,
        ) for start_char_ix, end_char_ix, label, kb_id in anns
    ]
    doc.ents = ents
    docbin.add(doc)

docbin.to_disk("./trn_docs.spacy")

Notice that I'm able to train a NER with this DocBin file and I do get meaningful results:

$ spacy train default01.cfg --paths.train ./trn_docs.spacy --paths.dev ./trn_docs.spacy --initialize.vectors en_core_web_lg --output ./my_ner
ℹ Using CPU

=========================== Initializing pipeline ===========================
[2021-09-28 18:20:26,609] [INFO] Set up nlp object from config
[2021-09-28 18:20:26,620] [INFO] Pipeline: ['tok2vec', 'ner']
[2021-09-28 18:20:26,624] [INFO] Created vocabulary
[2021-09-28 18:20:29,119] [INFO] Added vectors: en_core_web_lg
[2021-09-28 18:20:31,144] [INFO] Finished initializing nlp object
[2021-09-28 18:20:31,365] [INFO] Initialized pipeline components: ['tok2vec', 'ner']
✔ Initialized pipeline

============================= Training pipeline =============================
ℹ Pipeline: ['tok2vec', 'ner']
ℹ Initial learn rate: 0.001
E    #       LOSS TOK2VEC  LOSS NER  ENTS_F  ENTS_P  ENTS_R  SCORE
---  ------  ------------  --------  ------  ------  ------  ------
  0       0          0.00     32.79    0.00    0.00    0.00    0.00
200     200          0.32    318.50  100.00  100.00  100.00    1.00
400     400          0.00      0.00  100.00  100.00  100.00    1.00
600     600          0.00      0.00  100.00  100.00  100.00    1.00
800     800          0.00      0.00  100.00  100.00  100.00    1.00
1000    1000          0.00      0.00  100.00  100.00  100.00    1.00
1200    1200          0.00      0.00  100.00  100.00  100.00    1.00
1400    1400          0.00      0.00  100.00  100.00  100.00    1.00
1600    1600          0.00      0.00  100.00  100.00  100.00    1.00
1800    1800          0.00      0.00  100.00  100.00  100.00    1.00
✔ Saved pipeline to output directory
my_ner/model-last

The default01.cfg file contains the following:

[paths]
train = null
dev = null
vectors = null
init_tok2vec = null
raw_text = null

[system]
gpu_allocator = null
seed = 0

[nlp]
lang = "en"
pipeline = ["tok2vec","ner"]
batch_size = 1000
disabled = []
before_creation = null
after_creation = null
after_pipeline_creation = null
tokenizer = {"@tokenizers":"spacy.Tokenizer.v1"}

[components]

[components.ner]
factory = "ner"
incorrect_spans_key = null
moves = null
update_with_oracle_cut_size = 100

[components.ner.model]
@architectures = "spacy.TransitionBasedParser.v2"
state_type = "ner"
extra_state_tokens = false
hidden_width = 64
maxout_pieces = 2
use_upper = true
nO = null

[components.ner.model.tok2vec]
@architectures = "spacy.Tok2VecListener.v1"
width = ${components.tok2vec.model.encode.width}
upstream = "*"

[components.tok2vec]
factory = "tok2vec"

[components.tok2vec.model]
@architectures = "spacy.Tok2Vec.v2"

[components.tok2vec.model.embed]
@architectures = "spacy.MultiHashEmbed.v2"
width = ${components.tok2vec.model.encode.width}
attrs = ["NORM","PREFIX","SUFFIX","SHAPE"]
rows = [5000,2500,2500,2500]
include_static_vectors = true

[components.tok2vec.model.encode]
@architectures = "spacy.MaxoutWindowEncoder.v2"
width = 256
depth = 8
window_size = 1
maxout_pieces = 3

[corpora]

[corpora.dev]
@readers = "spacy.Corpus.v1"
path = ${paths.dev}
max_length = 0
gold_preproc = false
limit = 0
augmenter = null

[corpora.pretrain]
@readers = "spacy.Corpus.v1"
path = ${paths.train}
gold_preproc = false
max_length = 500
limit = 0

[corpora.train]
@readers = "spacy.Corpus.v1"
path = ${paths.train}
max_length = 0
gold_preproc = false
limit = 0
augmenter = null

[training]
dev_corpus = "corpora.dev"
train_corpus = "corpora.train"
seed = ${system.seed}
gpu_allocator = ${system.gpu_allocator}
dropout = 0.1
accumulate_gradient = 1
patience = 1600
max_epochs = 0
max_steps = 20000
eval_frequency = 200
frozen_components = []
annotating_components = []
before_to_disk = null

[training.batcher]
@batchers = "spacy.batch_by_words.v1"
discard_oversize = false
tolerance = 0.2
get_length = null

[training.batcher.size]
@schedules = "compounding.v1"
start = 100
stop = 1000
compound = 1.001
t = 0.0

[training.logger]
@loggers = "spacy.ConsoleLogger.v1"
progress_bar = false

[training.optimizer]
@optimizers = "Adam.v1"
beta1 = 0.9
beta2 = 0.999
L2_is_weight_decay = true
L2 = 0.01
grad_clip = 1.0
use_averages = false
eps = 0.00000001
learn_rate = 0.001

[training.score_weights]
ents_f = 1.0
ents_p = 0.0
ents_r = 0.0
ents_per_type = null

[pretraining]
max_epochs = 1000
dropout = 0.2
n_save_every = null
component = "tok2vec"
layer = ""
corpus = "corpora.pretrain"

[pretraining.batcher]
@batchers = "spacy.batch_by_words.v1"
size = 3000
discard_oversize = false
tolerance = 0.2
get_length = null

[pretraining.objective]
@architectures = "spacy.PretrainCharacters.v1"
maxout_pieces = 3
hidden_size = 300
n_characters = 4

[pretraining.optimizer]
@optimizers = "Adam.v1"
beta1 = 0.9
beta2 = 0.999
L2_is_weight_decay = true
L2 = 0.01
grad_clip = 1.0
use_averages = true
eps = 0.00000001
learn_rate = 0.001

[initialize]
vectors = "en_core_web_lg"
init_tok2vec = ${paths.init_tok2vec}
vocab_data = null
lookups = null
before_init = null
after_init = null

[initialize.components]

[initialize.tokenizer]

I then create the KB with the following code:

import spacy
from spacy.tokens import DocBin
from spacy.kb import KnowledgeBase
import pandas as pd
import numpy as np

nlp = spacy.load("./my_ner/model-last")
docs = list(DocBin().from_disk("./trn_docs.spacy").get_docs(nlp.vocab))
entity_ids = []
entity_strs = []
context_vectors = []
for doc in docs:
    for entity in doc.ents:
        entity_ids.append(entity.kb_id_)
        entity_strs.append(doc.text[entity.start_char:entity.end_char])
        context_vectors.append(entity.sent.vector)
entities = pd.DataFrame(
    {
        "EntityID": entity_ids,
        "EntityString": entity_strs,
    },
)
context_vectors = np.array(context_vectors)
vector_dim_cols = ["vd" + str(dim) for dim in range(context_vectors.shape[1])]
context_vectors = pd.DataFrame(context_vectors, columns=vector_dim_cols)
entities = pd.concat([entities, context_vectors], axis=1)

entity_freqs = entities.groupby("EntityID")["EntityID"].count()
entity_vectors = entities.groupby("EntityID")[vector_dim_cols].mean()

vectors_dim = nlp.vocab.vectors.shape[1]
kb = KnowledgeBase(vocab=nlp.vocab, entity_vector_length=vectors_dim)

for kbid, freq in entity_freqs.items():
    kb.add_entity(entity=kbid, entity_vector=entity_vectors.loc[kbid], freq=freq)

alias_counts = entities.groupby(["EntityString", "EntityID"])["EntityID"].count()
alias_probs = alias_counts / alias_counts.groupby("EntityString").sum()
for alias_string, group in alias_probs.groupby("EntityString"):
    group.reset_index(level="EntityString", drop=True, inplace=True)
    kb.add_alias(
        alias=alias_string,
        entities=group.index.values,
        probabilities=group.values,
    )

kb.to_disk("./my_kb.kb")

But then when I train the EL with this NER model and this KB I get 0 scores (but don't get the warning I was getting before):

$ python -m spacy train myconfig.cfg --paths.train ./trn_docs.spacy --paths.dev ./trn_docs.spacy --paths.kb ./my_kb.kb --paths.base_nlp ./my_ner/model-last --output ./my_el --code el_annotations.py
ℹ Using CPU

=========================== Initializing pipeline ===========================
[2021-09-29 15:54:33,415] [INFO] Set up nlp object from config
[2021-09-29 15:54:33,425] [INFO] Pipeline: ['sentencizer', 'ner', 'entity_linker']
[2021-09-29 15:54:33,430] [INFO] Created vocabulary
[2021-09-29 15:54:35,490] [INFO] Added vectors: ./my_ner/model-last
[2021-09-29 15:54:37,550] [INFO] Finished initializing nlp object
[2021-09-29 15:54:38,572] [INFO] Initialized pipeline components: ['entity_linker']
✔ Initialized pipeline

============================= Training pipeline =============================
ℹ Pipeline: ['sentencizer', 'ner', 'entity_linker']
ℹ Frozen components: ['sentencizer', 'ner']
ℹ Initial learn rate: 0.001
E    #       LOSS ENTIT...  SENTS_F  SENTS_P  SENTS_R  ENTS_F  ENTS_P  ENTS_R  NEL_MICRO_F  SCORE
---  ------  -------------  -------  -------  -------  ------  ------  ------  -----------  ------
  0       0           1.00   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
200     200          14.28   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
400     400           2.22   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
600     600           1.83   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
800     800           1.66   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
1000    1000           1.61   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
1200    1200           1.53   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
1400    1400           1.50   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
1600    1600           1.47   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
1800    1800           1.45   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
2000    2000           1.42   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
2200    2200           1.39   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
2400    2400           1.40   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
2600    2600           1.39   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
2800    2800           1.38   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
3000    3000           1.36   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
3200    3200           1.37   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
3400    3400           1.36   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
3600    3600           1.36   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
3800    3800           1.35   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
4000    4000           1.34   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
4200    4200           1.33   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
4400    4400           1.34   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
4600    4600           1.33   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
4800    4800           1.34   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
5000    5000           1.33   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
5200    5200           1.33   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
5400    5400           1.33   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
5600    5600           1.34   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
5800    5800           1.33   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
6000    6000           1.31   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
6200    6200           1.32   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
6400    6400           1.32   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
6600    6600           1.32   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
6800    6800           1.32   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
7000    7000           1.31   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
7200    7200           1.33   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
7400    7400           1.31   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
7600    7600           1.33   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
7800    7800           1.31   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
8000    8000           1.31   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
8200    8200           1.31   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
8400    8400           1.31   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
8600    8600           1.31   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
8800    8800           1.32   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
9000    9000           1.32   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
9200    9200           1.31   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
9400    9400           1.32   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
9600    9600           1.31   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
9800    9800           1.32   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
10000   10000           1.30   100.00   100.00   100.00    0.00    0.00    0.00         0.00    0.40
✔ Saved pipeline to output directory
my_el/model-last

[polm]

Sorry for the very delayed reply to this.

To come back to your initial goal:

there will be one and only one "candidate" kb_id for every entity in my training set.

Do you want the model to predict if the label is applicable to your candidate or not? Say something like "does Jaguar in this sentence refer to a cat or not" rather than "does Jaguar in this sentence refer to a cat or a car"?

If so that won't work with the EntityLinker. The way the model works is that for each possible candidate it generates a score, then takes the highest score, and that's the label. So if you have only one candidate it will just always use that label. The NIL label is only used if no candidates are available for an entity (or if the entity type is marked as to be ignored), not as a negative judgement. Adding score-based thresholding is something we intend to do eventually, but with the current architecture I think that would still only work with multiple candidates.

Also more generally I am not sure that task is generally considered entity linking, or it's not the standard construction. I'm not really sure what it would be called.

If I have misunderstood your goal, could you explain it in more detail, perhaps with a concrete example? (You sample data above didn't clarify things for me.)

---

Taking that into account and looking at your code, I suspect that what is happening is that you data currently has annotations (so you get no warning) but something is wonky with your candidate generator so you either always get predictions of NIL or something else is wrong. For debugging, try modifying your get_candidates function (whether you use the built-in one or not) to print what it's returning. Also maybe try looking at what happens with the list of candidates in the predict function of the EntityLinker.

[XBeg9]

Hi @alfredomg , were you able to find the root cause? I am getting the same issue, exact results as you are facing