quaterion.train.xbm.xbm_buffer module

class XbmBuffer(config: XbmConfig, embedding_size: int)

Bases: object

A buffer implementation to hold recent N embeddings and target values.

Inspired by https://github.com/msight-tech/research-xbm/blob/master/ret_benchmark/modeling/xbm.py

Parameters:

• config – Config class to configure XBM settings.

• embedding_size – Output dimension of EncoderHead configured for this model.

get() → Tuple[Tensor, LongTensor]

queue(embeddings: Tensor, targets: LongTensor) → None

Queue batch embeddings and targets in the buffer.

Parameters:

• embeddings – Output embeddings in the batch.

• targets – Target values in the batch.

property is_full: bool