Entity Linker - example config

[jkgenser]

The spacy quickstart config generators don't include any examples of entity linking. Does anyone have an example or recommended way to create entity linking configs for training? I can't seem to find any reference or examples, but I did see somewhere spacy staff mentioning to use the config for training models instead of manually training them.

Any help would be much appreciated!

https://spacy.io/usage/training#quickstart

[svlandeg]

Hi! We should definitely document this better and have this more accessible. What could be useful to you is this example project: https://github.com/explosion/projects/tree/v3/tutorials/nel_emerson. It shows how to train & run an entity_linker on some dummy data, and has an example config in the config subdir. And yes - we definitely recommend using this config for training models as there's a lot of parameters & settings that will be done correctly behind the scenes when using spacy train with a config file.

[jkgenser]

@svlandeg That's very helpful thank you! I ended up generating a config for my entity linker model after setting up my data. However I'm not getting an out of memory error with spacy trying to allocate ~18Gi of memory to my GPU when training. Traceback is reproduced below.

The line that is causing the GPU OOM is here: https://github.com/explosion/spaCy/blob/master/spacy/pipeline/entity_linker.py#L304

Spacy tries to initialize the entity linker on only 10 examples. The documents that is being passed to this are none longer than 5k characters. Later on in the initalization step it spacy.extract_spans is called on forward pass. This is there where the 18Gi memory allocation occurs. I've reproduced my config below as well.

Other posts on spacy memory allocation refer to batch size however, this operation the is OOMing seems to be unrelated to batch size as it has to do with the initialization step of the model itself before training even commences. Note: this also OOMs running on CPU on a machine with 32Gi memory.

I'm in the process of stepping through the spacy code to debug but not an expert in the library (yet) so any advice or notes on obvious gotchas are welcome :)

Update: I noticed if I change for eg in islice(get_examples(), 10): to try to init with 1 document instead of 1, the pipeline works. But then training immediately crashes. Perhaps the memory scaling with document size is not what I expected. I wonder if I need to chunk my documents to be much much smaller and having documents in the 5000 character range is too large for training this model? It's strange because I was able to train a model with NER only with no problem on GPU with similar and larger doc sizes. However for some reason the NER model is OOMing when it's part of the entity linking pipeline?

# function:
        self.model.initialize(
            X=doc_sample, Y=self.model.ops.asarray(vector_sample, dtype="float32")
        )

# docs in doc_sample
ipdb> for doc in doc_sample:print(len(doc))
3036
2658
4228
3330
1949
2505
1479
1737
3307
3772

Traceback (most recent call last):
  File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/spacy/__main__.py", line 4, in <module>
    setup_cli()
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/spacy/cli/_util.py", line 74, in setup_cli
    command(prog_name=COMMAND)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/click/core.py", line 1130, in __call__
    return self.main(*args, **kwargs)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/typer/core.py", line 778, in main
    return _main(
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/typer/core.py", line 216, in _main
    rv = self.invoke(ctx)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/click/core.py", line 1657, in invoke
    return _process_result(sub_ctx.command.invoke(sub_ctx))
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/click/core.py", line 1404, in invoke
    return ctx.invoke(self.callback, **ctx.params)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/click/core.py", line 760, in invoke
    return __callback(*args, **kwargs)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/typer/main.py", line 683, in wrapper
    return callback(**use_params)  # type: ignore
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/spacy/cli/train.py", line 45, in train_cli
    train(config_path, output_path, use_gpu=use_gpu, overrides=overrides)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/spacy/cli/train.py", line 72, in train
    nlp = init_nlp(config, use_gpu=use_gpu)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/spacy/training/initialize.py", line 84, in init_nlp    nlp.initialize(lambda: train_corpus(nlp), sgd=optimizer)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/spacy/language.py", line 1308, in initialize
    proc.initialize(get_examples, nlp=self, **p_settings)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/spacy/pipeline/entity_linker.py", line 304, in initialize
    self.model.initialize(
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/thinc/model.py", line 299, in initialize
    self.init(self, X=X, Y=Y)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/thinc/layers/chain.py", line 88, in init
    layer.initialize(X=curr_input, Y=Y)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/thinc/model.py", line 299, in initialize
    self.init(self, X=X, Y=Y)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/thinc/layers/chain.py", line 88, in init
    layer.initialize(X=curr_input, Y=Y)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/thinc/model.py", line 299, in initialize
    self.init(self, X=X, Y=Y)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/thinc/layers/chain.py", line 88, in init
    layer.initialize(X=curr_input, Y=Y)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/thinc/model.py", line 299, in initialize
    self.init(self, X=X, Y=Y)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/thinc/layers/chain.py", line 92, in init
    curr_input = layer.predict(curr_input)
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/thinc/model.py", line 315, in predict
    return self._func(self, X, is_train=False)[0]
  File "/home/j/.virtualenvs/oler-research/lib/python3.10/site-packages/spacy/ml/extract_spans.py", line 32, in forward
    Y = Ragged(X.dataXd[indices], spans.dataXd[:, 1] - spans.dataXd[:, 0])  # type: ignore[arg-type, index]
  File "cupy/_core/core.pyx", line 1519, in cupy._core.core._ndarray_base.__getitem__
  File "cupy/_core/_routines_indexing.pyx", line 44, in cupy._core._routines_indexing._ndarray_getitem
  File "cupy/_core/core.pyx", line 842, in cupy._core.core._ndarray_base.take
  File "cupy/_core/_routines_indexing.pyx", line 145, in cupy._core._routines_indexing._ndarray_take
  File "cupy/_core/_routines_indexing.pyx", line 802, in cupy._core._routines_indexing._take
  File "cupy/_core/core.pyx", line 136, in cupy._core.core.ndarray.__new__
  File "cupy/_core/core.pyx", line 224, in cupy._core.core._ndarray_base._init
  File "cupy/cuda/memory.pyx", line 742, in cupy.cuda.memory.alloc
  File "cupy/cuda/memory.pyx", line 1419, in cupy.cuda.memory.MemoryPool.malloc
  File "cupy/cuda/memory.pyx", line 1440, in cupy.cuda.memory.MemoryPool.malloc
  File "cupy/cuda/memory.pyx", line 1120, in cupy.cuda.memory.SingleDeviceMemoryPool.malloc
  File "cupy/cuda/memory.pyx", line 1141, in cupy.cuda.memory.SingleDeviceMemoryPool._malloc
  File "cupy/cuda/memory.pyx", line 1379, in cupy.cuda.memory.SingleDeviceMemoryPool._try_malloc
cupy.cuda.memory.OutOfMemoryError: Out of memory allocating 18,690,087,936 bytes (allocated so far: 2,050,890,240 bytes).

Config

[paths]
train = "/home/j/oler-research/data/nlp/ner/100/train.spacy"
dev = "/home/j/oler-research/data/nlp/ner/100/dev.spacy"
vectors = "en_core_web_lg"
init_tok2vec = null
kb = "/home/j/oler-research/data/nlp/ner/100/kb.spacy"

[system]
gpu_allocator = null
seed = 0

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

[components]

[components.entity_linker]
factory = "entity_linker"
candidates_batch_size = 1
# entity_vector_length = 64
entity_vector_length = 300
get_candidates = {"@misc":"spacy.CandidateGenerator.v1"}
get_candidates_batch = {"@misc":"spacy.CandidateBatchGenerator.v1"}
incl_context = true
incl_prior = true
labels_discard = []
n_sents = 0
overwrite = true
scorer = {"@scorers":"spacy.entity_linker_scorer.v1"}
threshold = null
use_gold_ents = true

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

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

[components.ner]
factory = "ner"
incorrect_spans_key = null
moves = null
scorer = {"@scorers":"spacy.ner_scorer.v1"}
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.sentencizer]
factory = "sentencizer"
overwrite = false
punct_chars = null
scorer = {"@scorers":"spacy.senter_scorer.v1"}

[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,1000,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.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 = 400
max_epochs = 0
max_steps = 20000
eval_frequency = 200
frozen_components = []
annotating_components = []
before_to_disk = null
before_update = 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]
sents_f = 0.20
sents_p = 0.0
sents_r = 0.0
ents_f = 0.40
ents_p = 0.0
ents_r = 0.0
ents_per_type = null
nel_micro_f = 0.40
nel_micro_r = null
nel_micro_p = null

[pretraining]

[initialize]
vectors = ${paths.vectors}
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]

[svlandeg]

Hey @jkgenser, can you tell me how large your KB is? The vector length is 300, right, and how many entities are in there? How many bytes does the file take up on disk?