Added mask_fill api and some examples in README
Summary:
1) This currently works only for single `<mask>` token as multi mask, we might have to look more into order of factorization.
2) This is currently only for single BPE token
Pull Request resolved:
@@ -133,3 +133,48 @@ class RobertaHubInterface(nn.Module):
assert len(doc) == aligned_feats.size(0)
doc.user_token_hooks['vector'] = lambda token: aligned_feats[token.i]
return doc
+
+ def fill_mask(self, masked_input: str, topk: int = 5):
+ masked_token = '<mask>'
+ assert masked_token in masked_input and masked_input.count(masked_token) == 1, \
+ "Please add one {0} token for the input, eg: 'He is a {0} guy'".format(masked_token)
+
+ text_spans = masked_input.split(masked_token)
+ text_spans_bpe = (' {0} '.format(masked_token)).join(
+ [self.bpe.encode(text_span.rstrip()) for text_span in text_spans]
+ ).strip()
+ tokens = self.task.source_dictionary.encode_line(
+ '<s> ' + text_spans_bpe,
+ append_eos=True,
+ )
+
+ masked_index = (tokens == self.task.mask_idx).nonzero()
+ if tokens.dim() == 1:
+ tokens = tokens.unsqueeze(0)
+
+ features, extra = self.model(
+ tokens.long().to(device=self.device),
+ features_only=False,
+ return_all_hiddens=False,
+ )
+ logits = features[0, masked_index, :].squeeze()
+ prob = logits.softmax(dim=0)
+ values, index = prob.topk(k=topk, dim=0)
+ topk_predicted_token_bpe = self.task.source_dictionary.string(index)
+
+ topk_filled_outputs = []
+ for index, predicted_token_bpe in enumerate(topk_predicted_token_bpe.split(' ')):
+ predicted_token = self.bpe.decode(predicted_token_bpe)
+ if " {0}".format(masked_token) in masked_input:
+ topk_filled_outputs.append((
+ masked_input.replace(
+ ' {0}'.format(masked_token), predicted_token
+ ),
+ values[index].item(),
+ ))
+ else:
+ topk_filled_outputs.append((
+ masked_input.replace(masked_token, predicted_token),
+ values[index].item(),
+ ))
+ return topk_filled_outputs
