"""BartphoTokenizer whose SentencePiece model is restricted to the 40K reduced vocab.

Stock BartphoTokenizer segments with the full 250K-piece SentencePiece model and
only afterwards maps pieces to the 40K fairseq dictionary: any piece outside the
dictionary becomes a single ``<unk>`` even when its characters are all known.
Restricting SentencePiece to the final vocabulary makes rare strings fall back to
smaller known pieces. The model was trained with this tokenizer, so inference must
use it too (load with ``AutoTokenizer.from_pretrained(..., trust_remote_code=True)``).
"""

from transformers import BartphoTokenizer as _BartphoTokenizer


class RestrictedBartphoTokenizer(_BartphoTokenizer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.restrict_sp_to_vocab()

    def restrict_sp_to_vocab(self):
        sp = self.sp_model
        sp.set_vocabulary([p for p in self.get_vocab() if sp.piece_to_id(p) != sp.unk_id()])

    def __setstate__(self, d):
        super().__setstate__(d)
        self.restrict_sp_to_vocab()