albert-khmer-small

import torch
from transformers import AlbertForMaskedLM, AlbertTokenizer
import sentencepiece as spm

# Load model and tokenizer
model = AlbertForMaskedLM.from_pretrained("seanghay/albert-khmer-small")
tokenizer = AlbertTokenizer.from_pretrained("seanghay/albert-khmer-small")
sp = spm.SentencePieceProcessor()
sp.load(tokenizer.vocab_file)

text = "ភ្នំពេញគឺជា[MASK]នៃប្រទេសកម្ពុជា។"
pieces = sp.encode_as_pieces(text)
ids = sp.encode_as_ids(text)
input_ids = torch.LongTensor([2] + ids + [3]).unsqueeze(0) # [CLS] + ids + [SEP]
attention_mask = torch.zeros_like(input_ids)

# Perform inference
with torch.no_grad():
    outputs = model(**inputs)
    logits = outputs.logits

# Locate the [MASK] token and extract predictions
mask_token_index = torch.where(inputs.input_ids == tokenizer.mask_token_id)[1]
mask_token_logits = logits[0, mask_token_index, :]

top_5_tokens = torch.topk(mask_token_logits, 5, dim=1).indices[0].tolist()

print(f"Original text: {text}")
print(f"Decoded input text embedding (should match original text): {sp.decode_ids(input_ids.squeeze().tolist())})
for i, token_id in enumerate(top_5_tokens):
    predicted_token = tokenizer.decode([token_id])
    print(f"{i + 1}. {text.replace('[MASK]', predicted_token)}")