Initial upload of custom Luganda embedding model architecture

8e1984a verified 2 days ago

8.13 kB

import os
import shutil
from typing import List, Optional, Tuple

import sentencepiece as spm
from transformers import PreTrainedTokenizer


class UnModelTokenizer(PreTrainedTokenizer):
    """SentencePiece tokenizer for the SAGE POND UnModel encoder.
    The special-token arguments are explicit constructor parameters. This is
    important for Transformers releases whose PythonBackend receives tokenizer
    configuration values through ``from_pretrained``: values such as
    ``bos_token`` must not also remain in ``kwargs`` when passed to ``super``.
    """

    vocab_files_names = {"vocab_file": "new_bpe.model"}
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(
        self,
        vocab_file,
        add_bos_token=False,
        add_eos_token=False,
        model_max_length=600,
        pad_token=None,
        bos_token=None,
        eos_token=None,
        unk_token=None,
        clean_up_tokenization_spaces=False,
        **kwargs,
    ):
        if not vocab_file or not os.path.isfile(vocab_file):
            raise FileNotFoundError(
                "The tokenizer requires the SentencePiece model file 'new_bpe.model'."
            )

        self.vocab_file = os.path.abspath(vocab_file)
        self.sp_model = spm.SentencePieceProcessor(model_file=self.vocab_file)
        self.add_bos_token = bool(add_bos_token)
        self.add_eos_token = bool(add_eos_token)

        vocab_size = self.sp_model.get_piece_size()
        if vocab_size != 40_000:
            raise ValueError(
                f"Expected a 40,000-piece SentencePiece model, got {vocab_size:,}."
            )

        def special_id(method_name, fallback):
            value = getattr(self.sp_model, method_name)()
            return value if value >= 0 else fallback

        # These fallbacks match the IDs used by the original training script.
        self._sp_pad_id = special_id("pad_id", 0)
        self._sp_bos_id = special_id("bos_id", 2)
        self._sp_eos_id = special_id("eos_id", 3)
        self._sp_unk_id = special_id("unk_id", 1)

        self._sp_special_tokens = {
            "pad_token": self.sp_model.id_to_piece(self._sp_pad_id),
            "bos_token": self.sp_model.id_to_piece(self._sp_bos_id),
            "eos_token": self.sp_model.id_to_piece(self._sp_eos_id),
            "unk_token": self.sp_model.id_to_piece(self._sp_unk_id),
        }

        for token_name, token_id in (
            ("pad_token", self._sp_pad_id),
            ("bos_token", self._sp_bos_id),
            ("eos_token", self._sp_eos_id),
            ("unk_token", self._sp_unk_id),
        ):
            if token_id < 0 or token_id >= vocab_size:
                raise ValueError(
                    f"{token_name} ID {token_id} is outside vocabulary size {vocab_size}."
                )

        # Validate values restored from tokenizer_config.json against the
        # SentencePiece model. Keeping the same strings guarantees that token
        # IDs remain aligned with the model's embedding table.
        supplied_tokens = {
            "pad_token": pad_token,
            "bos_token": bos_token,
            "eos_token": eos_token,
            "unk_token": unk_token,
        }
        resolved_tokens = {}
        for name, supplied in supplied_tokens.items():
            expected = self._sp_special_tokens[name]
            if supplied is None:
                resolved_tokens[name] = expected
                continue

            supplied_content = getattr(supplied, "content", str(supplied))
            if supplied_content != expected:
                raise ValueError(
                    f"Configured {name}={supplied_content!r} does not match the "
                    f"SentencePiece token {expected!r}. The tokenizer config and "
                    "new_bpe.model must come from the same training run."
                )
            resolved_tokens[name] = supplied

        super().__init__(
            pad_token=resolved_tokens["pad_token"],
            bos_token=resolved_tokens["bos_token"],
            eos_token=resolved_tokens["eos_token"],
            unk_token=resolved_tokens["unk_token"],
            model_max_length=model_max_length,
            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
            **kwargs,
        )

        # Ensure save_pretrained writes the dynamic AutoTokenizer mapping and
        # the options needed when the tokenizer is reconstructed from the Hub.
        self.init_kwargs["auto_map"] = {
            "AutoTokenizer": ["tokenization_un_model.UnModelTokenizer", None]
        }
        self.init_kwargs["tokenizer_class"] = self.__class__.__name__
        self.init_kwargs["add_bos_token"] = self.add_bos_token
        self.init_kwargs["add_eos_token"] = self.add_eos_token

    @property
    def vocab_size(self):
        return self.sp_model.get_piece_size()

    def get_vocab(self):
        vocab = {
            self.sp_model.id_to_piece(token_id): token_id
            for token_id in range(self.vocab_size)
        }
        vocab.update(self.added_tokens_encoder)
        return vocab

    def _tokenize(self, text, **kwargs):
        return self.sp_model.encode(text, out_type=str)

    def _convert_token_to_id(self, token):
        token_id = self.sp_model.piece_to_id(token)
        if token_id < 0 or token_id >= self.vocab_size:
            return self._sp_unk_id
        return token_id

    def _convert_id_to_token(self, index):
        if index < 0 or index >= self.vocab_size:
            return self.unk_token
        return self.sp_model.id_to_piece(index)

    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        return self.sp_model.decode_pieces(tokens)

    def build_inputs_with_special_tokens(
        self,
        token_ids_0: List[int],
        token_ids_1: Optional[List[int]] = None,
    ) -> List[int]:
        output = list(token_ids_0)
        if self.add_bos_token:
            output.insert(0, self.bos_token_id)
        if self.add_eos_token:
            output.append(self.eos_token_id)

        if token_ids_1 is not None:
            second = list(token_ids_1)
            if self.add_bos_token:
                second.insert(0, self.bos_token_id)
            if self.add_eos_token:
                second.append(self.eos_token_id)
            output.extend(second)
        return output

    def get_special_tokens_mask(
        self,
        token_ids_0: List[int],
        token_ids_1: Optional[List[int]] = None,
        already_has_special_tokens: bool = False,
    ) -> List[int]:
        if already_has_special_tokens:
            special_ids = set(self.all_special_ids)
            return [1 if token_id in special_ids else 0 for token_id in token_ids_0]

        mask = []
        if self.add_bos_token:
            mask.append(1)
        mask.extend([0] * len(token_ids_0))
        if self.add_eos_token:
            mask.append(1)

        if token_ids_1 is not None:
            if self.add_bos_token:
                mask.append(1)
            mask.extend([0] * len(token_ids_1))
            if self.add_eos_token:
                mask.append(1)
        return mask

    def create_token_type_ids_from_sequences(
        self,
        token_ids_0: List[int],
        token_ids_1: Optional[List[int]] = None,
    ) -> List[int]:
        combined = self.build_inputs_with_special_tokens(token_ids_0, token_ids_1)
        return [0] * len(combined)

    def save_vocabulary(
        self,
        save_directory,
        filename_prefix: Optional[str] = None,
    ) -> Tuple[str, ...]:
        if not os.path.isdir(save_directory):
            raise ValueError(f"Vocabulary save directory does not exist: {save_directory}")

        filename = (
            f"{filename_prefix}-new_bpe.model"
            if filename_prefix
            else "new_bpe.model"
        )
        destination = os.path.join(save_directory, filename)
        if os.path.abspath(self.vocab_file) != os.path.abspath(destination):
            shutil.copyfile(self.vocab_file, destination)
        return (destination,)


UnModelTokenizer.register_for_auto_class("AutoTokenizer")

	import os
	import shutil
	from typing import List, Optional, Tuple

	import sentencepiece as spm
	from transformers import PreTrainedTokenizer


	class UnModelTokenizer(PreTrainedTokenizer):
	"""SentencePiece tokenizer for the SAGE POND UnModel encoder.

	The special-token arguments are explicit constructor parameters. This is
	important for Transformers releases whose PythonBackend receives tokenizer
	configuration values through ``from_pretrained``: values such as
	``bos_token`` must not also remain in ``kwargs`` when passed to ``super``.
	"""

	vocab_files_names = {"vocab_file": "new_bpe.model"}
	model_input_names = ["input_ids", "attention_mask"]

	def __init__(
	self,
	vocab_file,
	add_bos_token=False,
	add_eos_token=False,
	model_max_length=600,
	pad_token=None,
	bos_token=None,
	eos_token=None,
	unk_token=None,
	clean_up_tokenization_spaces=False,
	**kwargs,
	):
	if not vocab_file or not os.path.isfile(vocab_file):
	raise FileNotFoundError(
	"The tokenizer requires the SentencePiece model file 'new_bpe.model'."
	)

	self.vocab_file = os.path.abspath(vocab_file)
	self.sp_model = spm.SentencePieceProcessor(model_file=self.vocab_file)
	self.add_bos_token = bool(add_bos_token)
	self.add_eos_token = bool(add_eos_token)

	vocab_size = self.sp_model.get_piece_size()
	if vocab_size != 40_000:
	raise ValueError(
	f"Expected a 40,000-piece SentencePiece model, got {vocab_size:,}."
	)

	def special_id(method_name, fallback):
	value = getattr(self.sp_model, method_name)()
	return value if value >= 0 else fallback

	# These fallbacks match the IDs used by the original training script.
	self._sp_pad_id = special_id("pad_id", 0)
	self._sp_bos_id = special_id("bos_id", 2)
	self._sp_eos_id = special_id("eos_id", 3)
	self._sp_unk_id = special_id("unk_id", 1)

	self._sp_special_tokens = {
	"pad_token": self.sp_model.id_to_piece(self._sp_pad_id),
	"bos_token": self.sp_model.id_to_piece(self._sp_bos_id),
	"eos_token": self.sp_model.id_to_piece(self._sp_eos_id),
	"unk_token": self.sp_model.id_to_piece(self._sp_unk_id),
	}

	for token_name, token_id in (
	("pad_token", self._sp_pad_id),
	("bos_token", self._sp_bos_id),
	("eos_token", self._sp_eos_id),
	("unk_token", self._sp_unk_id),
	):
	if token_id < 0 or token_id >= vocab_size:
	raise ValueError(
	f"{token_name} ID {token_id} is outside vocabulary size {vocab_size}."
	)

	# Validate values restored from tokenizer_config.json against the
	# SentencePiece model. Keeping the same strings guarantees that token
	# IDs remain aligned with the model's embedding table.
	supplied_tokens = {
	"pad_token": pad_token,
	"bos_token": bos_token,
	"eos_token": eos_token,
	"unk_token": unk_token,
	}
	resolved_tokens = {}
	for name, supplied in supplied_tokens.items():
	expected = self._sp_special_tokens[name]
	if supplied is None:
	resolved_tokens[name] = expected
	continue

	supplied_content = getattr(supplied, "content", str(supplied))
	if supplied_content != expected:
	raise ValueError(
	f"Configured {name}={supplied_content!r} does not match the "
	f"SentencePiece token {expected!r}. The tokenizer config and "
	"new_bpe.model must come from the same training run."
	)
	resolved_tokens[name] = supplied

	super().__init__(
	pad_token=resolved_tokens["pad_token"],
	bos_token=resolved_tokens["bos_token"],
	eos_token=resolved_tokens["eos_token"],
	unk_token=resolved_tokens["unk_token"],
	model_max_length=model_max_length,
	clean_up_tokenization_spaces=clean_up_tokenization_spaces,
	**kwargs,
	)

	# Ensure save_pretrained writes the dynamic AutoTokenizer mapping and
	# the options needed when the tokenizer is reconstructed from the Hub.
	self.init_kwargs["auto_map"] = {
	"AutoTokenizer": ["tokenization_un_model.UnModelTokenizer", None]
	}
	self.init_kwargs["tokenizer_class"] = self.__class__.__name__
	self.init_kwargs["add_bos_token"] = self.add_bos_token
	self.init_kwargs["add_eos_token"] = self.add_eos_token

	@property
	def vocab_size(self):
	return self.sp_model.get_piece_size()

	def get_vocab(self):
	vocab = {
	self.sp_model.id_to_piece(token_id): token_id
	for token_id in range(self.vocab_size)
	}
	vocab.update(self.added_tokens_encoder)
	return vocab

	def _tokenize(self, text, **kwargs):
	return self.sp_model.encode(text, out_type=str)

	def _convert_token_to_id(self, token):
	token_id = self.sp_model.piece_to_id(token)
	if token_id < 0 or token_id >= self.vocab_size:
	return self._sp_unk_id
	return token_id

	def _convert_id_to_token(self, index):
	if index < 0 or index >= self.vocab_size:
	return self.unk_token
	return self.sp_model.id_to_piece(index)

	def convert_tokens_to_string(self, tokens: List[str]) -> str:
	return self.sp_model.decode_pieces(tokens)

	def build_inputs_with_special_tokens(
	self,
	token_ids_0: List[int],
	token_ids_1: Optional[List[int]] = None,
	) -> List[int]:
	output = list(token_ids_0)
	if self.add_bos_token:
	output.insert(0, self.bos_token_id)
	if self.add_eos_token:
	output.append(self.eos_token_id)

	if token_ids_1 is not None:
	second = list(token_ids_1)
	if self.add_bos_token:
	second.insert(0, self.bos_token_id)
	if self.add_eos_token:
	second.append(self.eos_token_id)
	output.extend(second)
	return output

	def get_special_tokens_mask(
	self,
	token_ids_0: List[int],
	token_ids_1: Optional[List[int]] = None,
	already_has_special_tokens: bool = False,
	) -> List[int]:
	if already_has_special_tokens:
	special_ids = set(self.all_special_ids)
	return [1 if token_id in special_ids else 0 for token_id in token_ids_0]

	mask = []
	if self.add_bos_token:
	mask.append(1)
	mask.extend([0] * len(token_ids_0))
	if self.add_eos_token:
	mask.append(1)

	if token_ids_1 is not None:
	if self.add_bos_token:
	mask.append(1)
	mask.extend([0] * len(token_ids_1))
	if self.add_eos_token:
	mask.append(1)
	return mask

	def create_token_type_ids_from_sequences(
	self,
	token_ids_0: List[int],
	token_ids_1: Optional[List[int]] = None,
	) -> List[int]:
	combined = self.build_inputs_with_special_tokens(token_ids_0, token_ids_1)
	return [0] * len(combined)

	def save_vocabulary(
	self,
	save_directory,
	filename_prefix: Optional[str] = None,
	) -> Tuple[str, ...]:
	if not os.path.isdir(save_directory):
	raise ValueError(f"Vocabulary save directory does not exist: {save_directory}")

	filename = (
	f"{filename_prefix}-new_bpe.model"
	if filename_prefix
	else "new_bpe.model"
	)
	destination = os.path.join(save_directory, filename)
	if os.path.abspath(self.vocab_file) != os.path.abspath(destination):
	shutil.copyfile(self.vocab_file, destination)
	return (destination,)


	UnModelTokenizer.register_for_auto_class("AutoTokenizer")