kiwi.models.modules package

Submodules

kiwi.models.modules.attention module

class kiwi.models.modules.attention.Attention(scorer)

Bases: torch.nn.modules.module.Module

Generic Attention Implementation. Module computes a convex combination of a set of values based on the fit of their keys with a query.

forward(query, keys, values=None)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

set_mask(mask)

kiwi.models.modules.scorer module

class kiwi.models.modules.scorer.MLPScorer(query_size, key_size, layers=2, nonlinearity=<class 'torch.nn.modules.activation.Tanh'>)

Bases: kiwi.models.modules.scorer.Scorer

Implements a score function based on a Multilayer Perceptron.

forward(query, keys)

Computes Scores for each key given the query. :param query: FloatTensor batch x n :param keys: FloatTensor batch x seq_length x m

ret:

scores: FloatTensor batch x seq_length

class kiwi.models.modules.scorer.Scorer

Bases: torch.nn.modules.module.Module

Score function for Attention module.

forward(query, keys)

Computes Scores for each key given the query. :param query: FloatTensor batch x n :param keys: FloatTensor batch x seq_length x m

ret:

scores: FloatTensor batch x seq_length

Module contents