Tok2vec output tensor usage for search / similarity tasks

[darioprencipe]

Hello spaCy team,

I've recently trained ner and textcat components for a blank-sheet Italian spaCy language model that I use to both extract features (entities) from a very specific type of documents and to classify these documents. The 2 components share a tok2vec layer. I've done pretraining on a 1.5M documents corpus and the overall results for a small, compact model that is easily trainable in 2 hours on CPU (Apple M1, 8 cores and 8GB memory) are good enough. Long story short, no need for transformers (assuming transformers would perform better, given my document type and use case).

I've gone through both your docs and this nice primer, but I can't find a clear enough answer to the following questions:

Once trained my model, can I / does it make sense to use tok2vec output (i.e. Doc.tensor objects) as sort-of-contextual embeddings? Once I've learnt the distributions of tokens in my corpus, does it make sense to put together in some way tensors for all tokens in my Doc? So that I can, e.g., store these embeddings somewhere and use them to compute "contextual" (in the sense of my corpus) similarity across these documents?

Before venturing into testing tok2vec tensor-based similarity (e.g. storing these tensors somewhere and analysing them), I'd like to understand from you whether this makes sense or not, from both a theoretical and empirical perspective.

Here's my .cfg file setup so you can see how my training task looks like and understand what those tensors really mean.

[paths]
train = "training/data/train.spacy"
dev = "training/data/dev.spacy"
vectors = "it_core_news_lg"
init_tok2vec = "training/pretraining/output/model-last.bin"
raw_text = null

[system]
gpu_allocator = null
seed = 0

[nlp]
lang = "it"
pipeline = ["tok2vec","ner","textcat"]
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
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.textcat]
factory = "textcat"

[components.textcat.model]
@architectures = "spacy.TextCatEnsemble.v2"
nO = null

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

[components.textcat.model.linear_model]
@architectures = "spacy.TextCatBOW.v2"
exclusive_classes = true
ngram_size = 1
no_output_layer = false

[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.pretrain]
@readers = "spacy.JsonlCorpus.v1"
path = ${paths.raw_text}
min_length = 5
max_length = 500
limit = 0

[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 = 1600
max_epochs = -1
max_steps = 20000
eval_frequency = 500
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]
ents_f = 1.0
ents_p = 0.0
ents_r = 0.0
ents_per_type = null

[pretraining]
max_epochs = 30
dropout = 0.2
n_save_every = null
n_save_epoch = 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 = 1e-8
learn_rate = 0.001

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

[initialize.before_init]
@callbacks = "customize_tokenizer"

[initialize.components]

[initialize.components.ner]

[initialize.components.ner.labels]
@readers = "spacy.read_labels.v1"
path = "training/configs/labels/ner.json"

[initialize.components.textcat]

[initialize.components.textcat.labels]
@readers = "spacy.read_labels.v1"
path = "training/configs/labels/textcat.json"

[initialize.tokenizer]

Many many thanks in advance,

Dario