ctranslate2 4.7.0__cp314-cp314-macosx_11_0_arm64.whl

ctranslate2/specs/model_spec.py

@@ -0,0 +1,767 @@
+"""Specifications declare the expected variables layout of CTranslate2 models
+that do not load a computation graph. The model converter should make sure that
+each required variable of the specification is set.
+"""
+
+import abc
+import ctypes
+import json
+import os
+import shutil
+import struct
+
+from typing import Dict, List, Optional
+
+import numpy as np
+
+try:
+    import torch
+
+    torch_is_available = True
+except ImportError:
+    torch_is_available = False
+
+OPTIONAL = "__optional"
+CURRENT_BINARY_VERSION = 6
+
+ACCEPTED_MODEL_TYPES = (
+    "int8",
+    "int8_float32",
+    "int8_float16",
+    "int8_bfloat16",
+    "int16",
+    "float16",
+    "bfloat16",
+    "float32",
+)
+
+SKIP_CREATING_ALIAS = ("rotary_scaling_long_factor", "rotary_scaling_short_factor")
+
+
+def _join_scope(scope, name):
+    if not scope:
+        return name
+    return "%s/%s" % (scope, name)
+
+
+def _split_scope(scope):
+    return scope.split("/")
+
+
+def _parent_scope(scope):
+    keys = _split_scope(scope)
+    scope, attr = keys[:-1], keys[-1]
+    return "/".join(scope), attr
+
+
+def visit_spec(spec, fn, scope=""):
+    """Recursively visits a layer spec."""
+    for name, value in list(spec.__dict__.items()):
+        if name.startswith("_"):
+            continue
+        if isinstance(value, list):
+            for i, elem in enumerate(value):
+                visit_spec(elem, fn, scope=_join_scope(scope, "%s_%d" % (name, i)))
+        elif isinstance(value, LayerSpec):
+            visit_spec(value, fn, scope=_join_scope(scope, name))
+        else:
+            fn(spec, _join_scope(scope, name), value)
+
+
+def index_spec(spec, index):
+    if not index:
+        return spec
+    keys = _split_scope(index)
+    for key in keys:
+        try:
+            spec = getattr(spec, key)
+        except AttributeError:
+            attr, index = key.rsplit("_", 1)
+            spec = getattr(spec, attr)[int(index)]
+    return spec
+
+
+class FrozenMeta(type):
+    def __call__(self, *args, **kwargs):
+        instance = super().__call__(*args, **kwargs)
+        instance._frozen = True
+        return instance
+
+
+class FrozenAttr:
+    def __setattr__(self, key, value):
+        if hasattr(self, "_frozen") and not hasattr(self, key):
+            raise AttributeError("Attribute %s does not exist" % key)
+        super().__setattr__(key, value)
+
+
+class LayerSpec(FrozenAttr, metaclass=FrozenMeta):
+    """A layer specification declares the weights that should be set by the converters."""
+
+    def validate(self) -> None:
+        """Verify that the required weights are set.
+
+        Raises:
+          ValueError: If a required weight is not set in the specification.
+        """
+        unset_attributes = []
+
+        def _check(spec, name, value):
+            if value is None:
+                unset_attributes.append(name)
+                return
+
+            if isinstance(value, np.ndarray):
+                # float64 is not a supported type.
+                if value.dtype == np.float64:
+                    value = value.astype(np.float32)
+            elif isinstance(value, float):
+                value = np.dtype("float32").type(value)
+            elif isinstance(value, bool):
+                # Convert bool to an integer type.
+                value = np.dtype("int8").type(value)
+            elif isinstance(value, str):
+                if value != OPTIONAL:
+                    value = np.frombuffer(value.encode("utf-8"), dtype=np.int8)
+
+            if isinstance(value, np.ndarray) or isinstance(value, np.generic):
+                value = NumpyVariable(value)
+            elif torch_is_available and isinstance(value, torch.Tensor):
+                value = PyTorchVariable(value)
+
+            attr_name = _split_scope(name)[-1]
+            setattr(spec, attr_name, value)
+
+        self._visit(_check)
+
+        if unset_attributes:
+            raise ValueError(
+                "Some required model attributes are not set:\n\n%s"
+                % "\n".join(unset_attributes)
+            )
+
+    def variables(
+        self,
+        prefix: str = "",
+        ordered: bool = False,
+    ) -> Dict[str, np.ndarray]:
+        """Recursively returns the weights from this layer and its children.
+
+        Arguments:
+          prefix: Prefix to prepend to all variable names.
+          ordered: If set, an ordered list is returned instead.
+
+        Returns:
+          Dictionary mapping variables name to value.
+        """
+        var = {}
+
+        def _register_var(spec, name, value):
+            if isinstance(value, str) and value == OPTIONAL:
+                return
+            var[_join_scope(prefix, name)] = value
+
+        self._visit(_register_var)
+        if ordered:
+            return list(sorted(var.items(), key=lambda x: x[0]))
+        return var
+
+    def _alias_variables(self):
+        """Find duplicate variables in spec and create aliases."""
+        # When a variable is duplicated, keep the version that comes first in
+        # the alphabetical order and alias the others.
+        variables = self.variables(ordered=True)
+        for name, value in reversed(variables):
+            for other_name, other_value in variables:
+                if name == other_name:
+                    break
+                # Because variables can be transformed on load (e.g. transposed),
+                # we use an element-wise equality check.
+                scope, attr_name = _parent_scope(name)
+                if (
+                    not value.is_scalar()
+                    and value.equal(other_value)
+                    and attr_name not in SKIP_CREATING_ALIAS
+                ):
+                    # Replace variable value by the alias name.
+                    spec = index_spec(self, scope)
+                    setattr(spec, attr_name, other_name)
+                    break
+
+    def _quantize(self, quantization):
+        """Possibly quantizes the variable of the layer."""
+        if quantization is not None and quantization not in ACCEPTED_MODEL_TYPES:
+            raise ValueError(
+                "%s is not a valid quantization type. Accepted types are: %s"
+                % (quantization, ", ".join(ACCEPTED_MODEL_TYPES))
+            )
+
+        def _quantize(spec, name, value):
+            if not isinstance(value, Variable) or value.is_scalar():
+                return
+
+            key = _split_scope(name)[-1]
+            scale = None
+            is_quantizable = hasattr(spec, "%s_scale" % key)
+            is_convertible = value.dtype in ("float32", "float16", "bfloat16")
+
+            if is_quantizable:
+                if quantization == "int16":
+                    value = value.to("float32").numpy()
+                    # Represent the value with 10 bits so the multiplication is 20 bits
+                    # and 12 bits are left for accumulation.
+                    scale = np.float32(2**10 / np.amax(np.absolute(value)))
+                    value *= scale
+                    value = np.rint(value)
+                    value = np.clip(
+                        value, np.iinfo(np.int16).min, np.iinfo(np.int16).max
+                    )
+                    value = value.astype(np.int16)
+                    scale = NumpyVariable(scale)
+                    value = NumpyVariable(value)
+                elif quantization in (
+                    "int8",
+                    "int8_float32",
+                    "int8_float16",
+                    "int8_bfloat16",
+                ):
+                    value = value.to("float32").numpy()
+                    # For conv1d layer we need to reshape to 2D before calculating scale
+                    old_shape = None
+                    if len(value.shape) == 3:
+                        old_shape = value.shape
+                        value = value.reshape(value.shape[0], -1)
+                    amax = np.amax(np.absolute(value), axis=1)
+                    amax[amax == 0] = 127.0
+                    scale = 127.0 / amax
+                    value *= np.expand_dims(scale, 1)
+                    value = np.rint(value)
+                    value = value.astype(np.int8)
+                    # reshape back to old shape
+                    if old_shape:
+                        value = value.reshape(old_shape)
+                    scale = NumpyVariable(scale)
+                    value = NumpyVariable(value)
+                elif quantization in ("float16", "bfloat16", "float32"):
+                    value = value.to(quantization)
+
+            elif is_convertible:
+                if quantization in ("float16", "int8_float16"):
+                    value = value.to("float16")
+                elif quantization in ("bfloat16", "int8_bfloat16"):
+                    value = value.to("bfloat16")
+                elif quantization in ("float32", "int16", "int8_float32"):
+                    value = value.to("float32")
+
+            setattr(spec, key, value)
+            if scale is not None:
+                setattr(spec, "%s_scale" % key, scale)
+
+        self._visit(_quantize)
+
+    def optimize(self, quantization: Optional[str] = None) -> None:
+        """Recursively applies some optimizations to this layer:
+
+        * Alias variables with the same shape and value.
+        * Quantize weights.
+
+        Arguments:
+          quantization: Weight quantization scheme (possible values are: int8, int8_float32,
+            int8_float16, int8_bfloat16, int16, float16, bfloat16, float32).
+        """
+        self._alias_variables()
+        self._quantize(quantization)
+
+    def _visit(self, fn):
+        """Recursively visits this layer and its children."""
+        visit_spec(self, fn)
+
+
+def _dtype_to_type_id(object_dtype):
+    # Order should match the DataType enum in include/ctranslate2/types.h
+    dtypes = ("float32", "int8", "int16", "int32", "float16", "bfloat16")
+    try:
+        return dtypes.index(object_dtype)
+    except ValueError:
+        raise ValueError(
+            "%s is not in list of supported dtypes: %s"
+            % (object_dtype, ", ".join(dtypes))
+        )
+
+
+class ModelConfig(FrozenAttr, metaclass=FrozenMeta):
+    """Base class for model configurations."""
+
+    def __init__(self, **kwargs):
+        """Initializes the configuration with a set of parameters."""
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+    def to_dict(self):
+        """Returns the configuration as a dictionary."""
+        return {
+            key: value
+            for key, value in self.__dict__.items()
+            if not key.startswith("_")
+        }
+
+    def add_attribute(self, key, value):
+        self.__dict__[key] = value
+
+    def save_as_json(self, path):
+        """Saves the configuration as a JSON file."""
+        with open(path, "w", encoding="utf-8") as config_file:
+            json.dump(
+                self.to_dict(),
+                config_file,
+                indent=2,
+                sort_keys=True,
+            )
+            config_file.write("\n")
+
+
+class ModelSpec(LayerSpec):
+    """The top level layer specification."""
+
+    def __init__(self):
+        """Initializes the model specification."""
+        self._config = self.get_default_config()
+        self._files = {}
+
+    @property
+    def name(self):
+        """The name of the model specification."""
+        raise NotImplementedError()
+
+    @property
+    def revision(self):
+        """The model specification revision.
+
+        This value is incremented each time the weights layout of the model is
+        changed (e.g. a weight is renamed).
+        """
+        return 1
+
+    @property
+    def config(self):
+        """The model configuration."""
+        return self._config
+
+    def get_default_config(self):
+        """Returns the default configuration used by this model."""
+        return None
+
+    def register_file(self, path: str, filename: Optional[str] = None) -> None:
+        """Registers a file to be saved in the model directory."""
+        if not os.path.isfile(path):
+            raise ValueError("File %s does not exist" % path)
+        if filename is None:
+            filename = os.path.basename(path)
+        if filename in self._files:
+            raise ValueError("A file with name %s was already registered" % filename)
+        self._files[filename] = path
+
+    def save(self, output_dir: str) -> None:
+        """Saves this model on disk.
+
+        Arguments:
+          output_dir: Output directory where the model is saved.
+        """
+        self._serialize(os.path.join(output_dir, "model.bin"))
+        if self._config is not None:
+            self._config.save_as_json(os.path.join(output_dir, "config.json"))
+
+        for filename, path in self._files.items():
+            destination = os.path.join(output_dir, filename)
+            if os.path.exists(destination):
+                raise RuntimeError(
+                    "File %s already exists in the model directory" % destination
+                )
+            shutil.copy(path, destination)
+
+    def _serialize(self, path):
+        """Serializes the model variables."""
+        variables = []
+        aliases = []
+        for variable in self.variables(ordered=True):
+            if isinstance(variable[1], str):
+                aliases.append(variable)
+            else:
+                variables.append(variable)
+
+        with open(path, "wb") as model:
+
+            def _write_string(string):
+                model.write(struct.pack("H", len(string) + 1))
+                model.write(string.encode("utf-8"))
+                model.write(struct.pack("B", 0))
+
+            model.write(struct.pack("I", CURRENT_BINARY_VERSION))
+            _write_string(self.name)
+            model.write(struct.pack("I", self.revision))
+            model.write(struct.pack("I", len(variables)))
+            for name, value in variables:
+                _write_string(name)
+                model.write(struct.pack("B", len(value.shape)))
+                for dim in value.shape:
+                    model.write(struct.pack("I", dim))
+                model.write(struct.pack("B", _dtype_to_type_id(value.dtype)))
+                model.write(struct.pack("I", value.num_bytes()))
+                model.write(value.to_bytes())
+            model.write(struct.pack("I", len(aliases)))
+            for alias, variable_name in aliases:
+                _write_string(alias)
+                _write_string(variable_name)
+
+
+def _flatten_vocabularies(vocabularies):
+    for name, vocabulary in vocabularies.items():
+        if len(vocabulary) == 1:
+            yield name, vocabulary[0]
+        else:
+            for i, vocab in enumerate(vocabulary):
+                yield "%s_%d" % (name, i + 1), vocab
+
+
+class SequenceToSequenceModelConfig(ModelConfig):
+    """Configuration for sequence-to-sequence models."""
+
+    def __init__(
+        self,
+        unk_token: str = "<unk>",
+        bos_token: str = "<s>",
+        eos_token: str = "</s>",
+        decoder_start_token: Optional[str] = "<s>",
+        add_source_bos: bool = False,
+        add_source_eos: bool = False,
+        **kwargs,
+    ):
+        """Initializes the configuration for sequence-to-sequence models.
+
+        Args:
+          unk_token: The unknown token.
+          bos_token: The start of sentence token.
+          eos_token: The end of sentence token.
+          decoder_start_token: The decoder start token. If ``None``, the token should
+            be passed by the user in the target prefix.
+          add_source_bos: If ``True``, ``bos_token`` will be automatically added to
+            the source input.
+          add_source_eos: If ``True``, ``eos_token`` will be automatically added to
+            the source input.
+          **kwargs: Additional configuration.
+        """
+        super().__init__(
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            decoder_start_token=decoder_start_token,
+            add_source_bos=add_source_bos,
+            add_source_eos=add_source_eos,
+            **kwargs,
+        )
+
+
+class SequenceToSequenceModelSpec(ModelSpec):
+    """Base specification for sequence to sequence models."""
+
+    def __init__(self):
+        """Initializes a sequence to sequence model specification."""
+        super().__init__()
+        self._vocabularies = {
+            "source": [],
+            "target": [],
+        }
+
+    def get_default_config(self):
+        return SequenceToSequenceModelConfig()
+
+    @abc.abstractmethod
+    def get_source_vocabulary_size(self):
+        """Returns the source vocabulary size expected by the model."""
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_target_vocabulary_size(self):
+        """Returns the target vocabulary size expected by the model."""
+        raise NotImplementedError()
+
+    def register_source_vocabulary(self, tokens: List[str]) -> None:
+        """Registers a source vocabulary of tokens.
+
+        Arguments:
+          tokens: List of source tokens.
+        """
+        self._vocabularies["source"].append(tokens)
+
+    def register_target_vocabulary(self, tokens: List[str]) -> None:
+        """Registers a target vocabulary of tokens.
+
+        Arguments:
+          tokens: List of target tokens.
+        """
+        self._vocabularies["target"].append(tokens)
+
+    def register_vocabulary_mapping(self, path: str) -> None:
+        """Registers a vocabulary mapping file.
+
+        Arguments:
+          path: Path to the vocabulary mapping file.
+        """
+        self.register_file(path, "vmap.txt")
+
+    def validate(self) -> None:
+        super().validate()
+
+        # Check that vocabularies are registered and have the correct size.
+        vocabulary_sizes = {
+            "source": self.get_source_vocabulary_size(),
+            "target": self.get_target_vocabulary_size(),
+        }
+
+        for name, sizes in vocabulary_sizes.items():
+            if not isinstance(sizes, list):
+                sizes = [sizes]
+            vocabularies = self._vocabularies[name]
+            if len(vocabularies) != len(sizes):
+                raise ValueError(
+                    "Incorrect number of %s vocabularies: %d registered, but expected %d"
+                    % (name, len(vocabularies), len(sizes))
+                )
+            for i, (vocabulary, expected_size) in enumerate(zip(vocabularies, sizes)):
+                if len(vocabulary) != expected_size:
+                    raise ValueError(
+                        "%s vocabulary %d has size %d but the model expected a vocabulary "
+                        "of size %d"
+                        % (name.capitalize(), i, len(vocabulary), expected_size)
+                    )
+
+    def save(self, output_dir: str) -> None:
+        # Save the vocabularies.
+        vocabularies = dict(_flatten_vocabularies(self._vocabularies))
+        all_vocabularies = list(vocabularies.values())
+        if all(vocabulary == all_vocabularies[0] for vocabulary in all_vocabularies):
+            vocabularies = {"shared": all_vocabularies[0]}
+
+        for name, tokens in vocabularies.items():
+            _save_vocabulary(output_dir, "%s_vocabulary" % name, tokens)
+
+        # Save the rest of the model.
+        super().save(output_dir)
+
+
+class LanguageModelConfig(ModelConfig):
+    """Configuration for language models."""
+
+    def __init__(
+        self,
+        unk_token: str = "<unk>",
+        bos_token: str = "<s>",
+        eos_token: str = "</s>",
+        **kwargs,
+    ):
+        """Initializes the configuration for language models.
+
+        Args:
+          unk_token: The unknown token.
+          bos_token: The start of sentence token.
+          eos_token: The end of sentence token.
+          **kwargs: Additional configuration.
+        """
+        super().__init__(
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            **kwargs,
+        )
+
+
+class LanguageModelSpec(ModelSpec):
+    """Base specification for language models."""
+
+    def __init__(self):
+        """Initializes a language model specification."""
+        super().__init__()
+        self._vocabulary = []
+
+    def get_default_config(self):
+        return LanguageModelConfig()
+
+    @abc.abstractmethod
+    def get_vocabulary_size(self):
+        """Returns the vocabulary size expected by the model."""
+        raise NotImplementedError()
+
+    def register_vocabulary(self, tokens: List[str]) -> None:
+        """Registers the vocabulary of tokens.
+
+        Arguments:
+          tokens: List of tokens.
+        """
+        self._vocabulary = list(tokens)
+
+    def validate(self) -> None:
+        super().validate()
+
+        expected_vocabulary_size = self.get_vocabulary_size()
+        if len(self._vocabulary) != expected_vocabulary_size:
+            raise ValueError(
+                "Vocabulary has size %d but the model expected a vocabulary of size %d"
+                % (len(self._vocabulary), expected_vocabulary_size)
+            )
+
+    def save(self, output_dir: str) -> None:
+        # Save the vocabulary.
+        _save_vocabulary(output_dir, "vocabulary", self._vocabulary)
+
+        # Save the rest of the model.
+        super().save(output_dir)
+
+
+def _save_vocabulary(output_dir, name, tokens):
+    vocabulary_path = os.path.join(output_dir, "%s.json" % name)
+
+    with open(vocabulary_path, "w", encoding="utf-8") as vocabulary_file:
+        json.dump(tokens, vocabulary_file, indent=2)
+
+
+class Variable(abc.ABC):
+    """Abstract base class for model variables."""
+
+    @property
+    @abc.abstractmethod
+    def shape(self) -> List[int]:
+        raise NotImplementedError()
+
+    def is_scalar(self) -> bool:
+        return len(self.shape) == 0
+
+    @property
+    @abc.abstractmethod
+    def dtype(self) -> str:
+        raise NotImplementedError()
+
+    def to(self, dtype: str) -> "Variable":
+        if dtype == self.dtype:
+            return self
+        return self._to(dtype)
+
+    @abc.abstractmethod
+    def numpy(self) -> np.ndarray:
+        raise NotImplementedError()
+
+    def equal(self, other) -> bool:
+        return type(self) is type(other) and self._equal(other)
+
+    @abc.abstractmethod
+    def num_bytes(self) -> int:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def to_bytes(self) -> bytes:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def _to(self, dtype: str) -> "Variable":
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def _equal(self, other) -> bool:
+        raise NotImplementedError()
+
+
+class NumpyVariable(Variable):
+    """Model variable as a Numpy array."""
+
+    def __init__(self, array):
+        self.array = array
+
+    @property
+    def shape(self) -> List[int]:
+        return self.array.shape
+
+    @property
+    def dtype(self) -> str:
+        return self.array.dtype.name
+
+    def numpy(self) -> np.ndarray:
+        return self.array
+
+    def num_bytes(self) -> int:
+        return self.array.nbytes
+
+    def to_bytes(self) -> bytes:
+        return self.array.tobytes()
+
+    def _to(self, dtype: str) -> Variable:
+        if dtype == "bfloat16":
+            if not torch_is_available:
+                raise RuntimeError(
+                    "Converting to bfloat16 requires torch to be installed"
+                )
+            return PyTorchVariable.from_numpy(self.array).to(dtype)
+
+        dtype = np.dtype(dtype)
+        self.array = self.array.astype(dtype)
+        return self
+
+    def _equal(self, other) -> bool:
+        a = self.array
+        b = other.array
+        return a is b or (
+            a.dtype == b.dtype
+            and a.shape == b.shape
+            and a.flat[0] == b.flat[0]
+            and np.array_equal(a, b)
+        )
+
+
+class PyTorchVariable(Variable):
+    """Model variable as a PyTorch tensor."""
+
+    def __init__(self, tensor):
+        if isinstance(tensor, torch.nn.Parameter):
+            tensor = tensor.data
+
+        self.tensor = tensor.contiguous()
+
+    @classmethod
+    def from_numpy(cls, array):
+        tensor = torch.from_numpy(array)
+        return cls(tensor)
+
+    @property
+    def shape(self) -> List[int]:
+        return list(self.tensor.shape)
+
+    @property
+    def dtype(self) -> str:
+        return str(self.tensor.dtype).replace("torch.", "")
+
+    def numpy(self) -> np.ndarray:
+        return self.tensor.detach().numpy()
+
+    def num_bytes(self) -> int:
+        return self.tensor.numel() * self.tensor.element_size()
+
+    def to_bytes(self) -> bytes:
+        max_size = 2**31 - 1
+        num_bytes = self.num_bytes()
+        output = b""
+        offset = 0
+        while num_bytes > 0:
+            chunk_size = max_size if num_bytes > max_size else num_bytes
+            chunk = ctypes.string_at(self.tensor.data_ptr() + offset, chunk_size)
+            output += chunk
+            offset += chunk_size
+            num_bytes -= chunk_size
+        return output
+
+    def _to(self, dtype: str) -> Variable:
+        dtype = getattr(torch, dtype)
+        self.tensor = self.tensor.to(dtype)
+        return self
+
+    def _equal(self, other) -> bool:
+        a = self.tensor
+        b = other.tensor
+        return a is b or (a.dtype == b.dtype and torch.equal(a, b))