def create_and_check_context_encoder( self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels ): model = DPRContextEncoder(config=config) model.to(torch_device) model.eval() result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids) result = model(input_ids, token_type_ids=token_type_ids) result = model(input_ids) self.parent.assertEqual(result.pooler_output.shape, (self.batch_size, self.projection_dim or self.hidden_size))
def create_and_check_dpr_context_encoder(self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels): model = DPRContextEncoder(config=config) model.to(torch_device) model.eval() embeddings = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids)[0] embeddings = model(input_ids, token_type_ids=token_type_ids)[0] embeddings = model(input_ids)[0] result = { "embeddings": embeddings, } self.parent.assertListEqual( list(result["embeddings"].size()), [self.batch_size, self.projection_dim or self.hidden_size])