Update README.md

README.md

@@ -34,6 +34,36 @@ tokenizer = AutoTokenizer.from_pretrained("Yaxin/roberta-large-ernie2-skep-en")
 model = AutoModel.from_pretrained("Yaxin/roberta-large-ernie2-skep-en")
 ```
 
+```
+#!/usr/bin/env python
+#encoding: utf-8
+import torch
+from transformers import RobertaTokenizer, RobertaForMaskedLM
+
+tokenizer = RobertaTokenizer.from_pretrained('Yaxin/roberta-large-ernie2-skep-en')
+
+input_tx = "<s> He like play with student, so he became a <mask> after graduation </s>"
+# input_tx = "<s> He is a <mask> and likes to get along with his students </s>"
+
+tokenized_text = tokenizer.tokenize(input_tx)
+indexed_tokens = tokenizer.convert_tokens_to_ids(tokenized_text)
+
+tokens_tensor = torch.tensor([indexed_tokens])
+segments_tensors = torch.tensor([[0] * len(tokenized_text)])
+
+model = RobertaForMaskedLM.from_pretrained('Yaxin/roberta-large-ernie2-skep-en')
+model.eval()
+
+with torch.no_grad():
+    outputs = model(tokens_tensor, token_type_ids=segments_tensors)
+    predictions = outputs[0]
+
+predicted_index = [torch.argmax(predictions[0, i]).item() for i in range(0, (len(tokenized_text) - 1))]
+predicted_token = [tokenizer.convert_ids_to_tokens([predicted_index[x]])[0] for x in
+                   range(1, (len(tokenized_text) - 1))]
+
+print('Predicted token is:', predicted_token)
+```
 ## Citation
 
 ```bibtex