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

class capreolus.reranker.TK.PositionalEncoding(d_model, dropout=0.1, max_len=5000)

Bases: torch.nn.Module

forward(x)

class capreolus.reranker.TK.TK_class(extractor, config)

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)

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)

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

get_embedding(toks)

Overrides KNRM_Class’s get_embedding to return contextualized word embeddings

forward(doctoks, querytoks, query_idf)

class capreolus.reranker.TK.TK(config=None, provide=None, share_dependency_objects=False, build=True)

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'

dependencies

config_spec

build_model()

score(d)

test(d)