primeqa.ir.dense.dpr_top.dpr.dataloader_biencoder.BiEncoderLoader

class primeqa.ir.dense.dpr_top.dpr.dataloader_biencoder.BiEncoderLoader(hypers: primeqa.ir.dense.dpr_top.dpr.biencoder_hypers.BiEncoderHypers, per_gpu_batch_size: int, qry_tokenizer, ctx_tokenizer, data_dir, positive_pid_file, *, files_per_dataloader=1, checkpoint_info=None)

Bases: primeqa.ir.dense.dpr_top.dataloader.distloader_base.MultiFileLoader

Methods

all_batches
batch_dict	param batch input_ids_q, attention_mask_q, input_ids_c, attention_mask_c, positive_indices
display_batch
get_checkpoint_info
get_dataloader
quick_test
reset

batch_dict(batch)

Parameters

batch – input_ids_q, attention_mask_q, input_ids_c, attention_mask_c, positive_indices

Returns