def bert_forward(input_ids, attention_mask, token_type_ids):
input_model = {
'input_ids': input_ids.long(),
'attention_mask': attention_mask.long(),
'token_type_ids': token_type_ids.long(),
}
output_model = model(**input_model)[0]
return output_model
def roberta_forward(input_ids, attention_mask):
input_model = {
'input_ids': input_ids.long(),
'attention_mask': attention_mask.long(),
}
output_model = model(**input_model)[0]
return output_model