capreolus.reranker.TK

Module Contents

Classes

PositionalEncoding

TK_class

Adapted from https://github.com/sebastian-hofstaetter/transformer-kernel-ranking/blob/master/matchmaker/models/tk.py

TK

Sebastian Hofstätter, Markus Zlabinger, and Allan Hanbury. 2019. TU Wien @ TREC Deep Learning '19 -- Simple Contextualization for Re-ranking. In TREC '19.

Attributes

logger

capreolus.reranker.TK.logger[source]
class capreolus.reranker.TK.PositionalEncoding(d_model, dropout=0.1, max_len=5000)[source]

Bases: torch.nn.Module

forward(x)[source]
class capreolus.reranker.TK.TK_class(extractor, config)[source]

Bases: torch.nn.Module

Adapted from https://github.com/sebastian-hofstaetter/transformer-kernel-ranking/blob/master/matchmaker/models/tk.py TK is a neural IR model - a fusion between transformer contextualization & kernel-based scoring -> uses 1 transformer block to contextualize embeddings -> soft-histogram kernels to score interactions

get_mu_matrix(extractor)[source]

Returns a matrix of mu values that can be directly subtracted from the cosine matrix. This is the matrix mu in equation 5 in the paper (https://arxiv.org/pdf/2002.01854.pdf)

get_mask(embedding)[source]

Gets a mask of shape (seq_len, seq_len). This is an additive mask, hence masked elements should be -inf

get_embedding(toks)[source]

Overrides KNRM_Class’s get_embedding to return contextualized word embeddings

forward(doctoks, querytoks, query_idf)[source]
class capreolus.reranker.TK.TK(config=None, provide=None, share_dependency_objects=False, build=True)[source]

Bases: capreolus.reranker.Reranker

Sebastian Hofstätter, Markus Zlabinger, and Allan Hanbury. 2019. TU Wien @ TREC Deep Learning ‘19 – Simple Contextualization for Re-ranking. In TREC ‘19.

module_name = 'TK'[source]
dependencies[source]
config_spec[source]
build_model()[source]
score(d)[source]
test(d)[source]