Update README.md

README.md

@@ -22,24 +22,16 @@ from transformers import RobertaTokenizerFast, RobertaForMaskedLM
 import torch
 # Load the tokenizer
 tokenizer = RobertaTokenizerFast.from_pretrained('minhtoan/roberta-masked-lm-vietnamese-nom')
-
 # Load the model
 model = RobertaForMaskedLM.from_pretrained('minhtoan/roberta-masked-lm-vietnamese-nom')
 
-# Example input sentence with a masked token
-input_sentence = '想払𨀐' + '<mask>'
-
-# Tokenize the input sentence
-input_tokens = tokenizer.encode(input_sentence, return_tensors='pt')
-
-# Generate predictions
-with torch.no_grad():
-    outputs = model(input_tokens)
-    predictions = outputs.logits.argmax(dim=-1)
+text = '<mask>如㗂䳽𠖤戈'
+inputs = tokenizer(text, return_tensors="pt")
+mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]
 
-# Decode and print the predicted word
-predicted_word = tokenizer.decode(predictions[0, -1].item())
-print("Predicted word:", predicted_word)
+logits = model(**inputs).logits
+mask_token_logits = logits[0, mask_token_index, :]
+print("Predicted word:",  tokenizer.decode(mask_token_logits[0].argmax()))
 ~~~~