"""Run locally: python example.py /path/to/release (or a Hub repo ID)."""
import sys
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM
checkpoint = sys.argv[1] if len(sys.argv) > 1 else 'Abdelkareem/arabic-berthash-v2'
tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
model = AutoModelForMaskedLM.from_pretrained(checkpoint, trust_remote_code=True).eval()
batch = tokenizer('اللغة العربية هي لغة [MASK] الكريم',return_tensors='pt',truncation=True,max_length=1024)
mask = batch['input_ids'].eq(tokenizer.mask_token_id)
assert mask.sum().item() == 1
with torch.inference_mode():
    logits = model(**batch).logits[mask][0]
    probabilities, ids = logits.softmax(-1).topk(5)
for token, probability in zip(ids.tolist(),probabilities.tolist()):
    filled=batch['input_ids'][0].clone();filled[mask[0]]=token
    print(tokenizer.decode(filled,skip_special_tokens=True),round(probability,5))