Update README.md

README.md

@@ -34,19 +34,30 @@ import torch
 tokenizer = BertTokenizer.from_pretrained('suayptalha/medBERT-base')
 model = BertForMaskedLM.from_pretrained('suayptalha/medBERT-base').to("cuda")
 
-input_text = "The patient was diagnosed with gastric cancer after a thorough examination."
-masked_text = input_text.replace("gastric cancer", tokenizer.mask_token)
-
-inputs = tokenizer(masked_text, return_tensors='pt').to("cuda")
+input_text = "Response to neoadjuvant chemotherapy best predicts survival [MASK] curative resection of gastric cancer."
+inputs = tokenizer(input_text, return_tensors='pt').to("cuda")
 
 outputs = model(**inputs)
 
-predicted_token_id = torch.argmax(outputs.logits, dim=-1)
+masked_index = (inputs['input_ids'][0] == tokenizer.mask_token_id).nonzero(as_tuple=True)[0].item()
+
+top_k = 5
+logits = outputs.logits[0, masked_index]
+top_k_ids = torch.topk(logits, k=top_k).indices.tolist()
+top_k_tokens = tokenizer.convert_ids_to_tokens(top_k_ids)
 
-predicted_token = tokenizer.decode(predicted_token_id[0, inputs['input_ids'].shape[1] - 1])
-print(predicted_token)
+print("Top 5 prediction:")
+for i, token in enumerate(top_k_tokens):
+    print(f"{i + 1}: {token}")
 '''
 
+_Top 5 prediction:_
+_1: from_
+_2: of_
+_3: after_
+_4: by_
+_5: through_
+
 ### **Fine-tuning the Model**
 
 To fine-tune the **medBERT-base** model on your own medical dataset, follow these steps: