ctranslate2 4.7.0__cp314-cp314-macosx_11_0_arm64.whl

ctranslate2/converters/fairseq.py

@@ -0,0 +1,347 @@
+import argparse
+import os
+
+from typing import Optional
+
+from ctranslate2.converters import utils
+from ctranslate2.converters.converter import Converter
+from ctranslate2.specs import common_spec, transformer_spec
+
+_SUPPORTED_MODELS = {
+    "bart",
+    "multilingual_transformer",
+    "transformer",
+    "transformer_align",
+    "transformer_lm",
+}
+
+
+_SUPPORTED_ACTIVATIONS = {
+    "gelu": common_spec.Activation.GELU,
+    "gelu_accurate": common_spec.Activation.GELUTanh,
+    "gelu_fast": common_spec.Activation.GELUTanh,
+    "relu": common_spec.Activation.RELU,
+    "swish": common_spec.Activation.SWISH,
+}
+
+
+def _get_model_spec(args):
+    import fairseq
+
+    activation_fn = getattr(args, "activation_fn", "relu")
+    model_name = fairseq.models.ARCH_MODEL_NAME_REGISTRY[args.arch]
+
+    check = utils.ConfigurationChecker()
+    check(
+        model_name in _SUPPORTED_MODELS,
+        "Model '%s' used by architecture '%s' is not supported (supported models are: %s)"
+        % (model_name, args.arch, ", ".join(_SUPPORTED_MODELS)),
+    )
+    check.validate()
+    check(
+        activation_fn in _SUPPORTED_ACTIVATIONS,
+        "Option --activation-fn %s is not supported (supported activations are: %s)"
+        % (activation_fn, ", ".join(_SUPPORTED_ACTIVATIONS.keys())),
+    )
+    check(
+        not getattr(args, "no_token_positional_embeddings", False),
+        "Option --no-token-positional-embeddings is not supported",
+    )
+    check(
+        not getattr(args, "lang_tok_replacing_bos_eos", False),
+        "Option --lang-tok-replacing-bos-eos is not supported",
+    )
+
+    if model_name == "transformer_lm":
+        check(
+            not args.character_embeddings,
+            "Option --character-embeddings is not supported",
+        )
+        check(
+            not args.adaptive_input,
+            "Option --adaptive-input is not supported",
+        )
+        check.validate()
+
+        return transformer_spec.TransformerDecoderModelSpec.from_config(
+            args.decoder_layers,
+            args.decoder_attention_heads,
+            pre_norm=args.decoder_normalize_before,
+            activation=_SUPPORTED_ACTIVATIONS[activation_fn],
+            layernorm_embedding=getattr(args, "layernorm_embedding", False),
+            no_final_norm=args.no_decoder_final_norm,
+            project_in_out=args.decoder_input_dim != args.decoder_embed_dim,
+        )
+
+    else:
+        check(
+            args.encoder_normalize_before == args.decoder_normalize_before,
+            "Options --encoder-normalize-before and --decoder-normalize-before "
+            "must have the same value",
+        )
+        check(
+            args.encoder_attention_heads == args.decoder_attention_heads,
+            "Options --encoder-attention-heads and --decoder-attention-heads "
+            "must have the same value",
+        )
+        check.validate()
+
+        return transformer_spec.TransformerSpec.from_config(
+            (args.encoder_layers, args.decoder_layers),
+            args.encoder_attention_heads,
+            pre_norm=args.encoder_normalize_before,
+            activation=_SUPPORTED_ACTIVATIONS[activation_fn],
+            alignment_layer=getattr(args, "alignment_layer", -1),
+            alignment_heads=getattr(args, "alignment_heads", 0),
+            layernorm_embedding=getattr(args, "layernorm_embedding", False),
+        )
+
+
+def _get_vocab(dictionary):
+    return ["<blank>" if token == "<pad>" else token for token in dictionary.symbols]
+
+
+class FairseqConverter(Converter):
+    """Converts models trained with Fairseq."""
+
+    def __init__(
+        self,
+        model_path: str,
+        data_dir: str,
+        source_lang: Optional[str] = None,
+        target_lang: Optional[str] = None,
+        fixed_dictionary: Optional[str] = None,
+        no_default_special_tokens: bool = False,
+        user_dir: Optional[str] = None,
+    ):
+        """Initializes the Fairseq converter.
+
+        Arguments:
+          model_path: Path to the Fairseq PyTorch model (.pt file).
+          data_dir: Path to the Fairseq data directory containing vocabulary files.
+          source_lang: Source language (may be required if not declared in the model).
+          target_lang: Target language (may be required if not declared in the model).
+          fixed_dictionary: Path to the fixed dictionary for multilingual models.
+          no_default_special_tokens: Require all special tokens to be provided by the user
+            (e.g. encoder end token, decoder start token).
+          user_dir: Path to the user directory containing custom extensions.
+        """
+        self._model_path = model_path
+        self._data_dir = data_dir
+        self._fixed_dictionary = fixed_dictionary
+        self._source_lang = source_lang
+        self._target_lang = target_lang
+        self._no_default_special_tokens = no_default_special_tokens
+        self._user_dir = user_dir
+
+    def _load(self):
+        import fairseq
+        import torch
+
+        from fairseq import checkpoint_utils
+
+        if self._user_dir:
+            from fairseq.utils import import_user_module
+
+            import_user_module(argparse.Namespace(user_dir=self._user_dir))
+
+        with torch.no_grad():
+            checkpoint = torch.load(
+                self._model_path, map_location=torch.device("cpu"), weights_only=False
+            )
+            args = checkpoint["args"] or checkpoint["cfg"]["model"]
+
+            args.data = self._data_dir
+            if self._fixed_dictionary is not None:
+                args.fixed_dictionary = self._fixed_dictionary
+            if hasattr(args, "lang_dict") and args.lang_dict:
+                args.lang_dict = os.path.join(
+                    self._data_dir, os.path.basename(args.lang_dict)
+                )
+
+            if self._source_lang is not None:
+                args.source_lang = self._source_lang
+
+            if self._target_lang is not None:
+                args.target_lang = self._target_lang
+
+            spec = _get_model_spec(args)
+
+            task = fairseq.tasks.setup_task(args)
+            model = fairseq.models.build_model(args, task)
+            model.eval()
+            model.load_state_dict(checkpoint["model"])
+
+            if isinstance(spec, transformer_spec.TransformerDecoderModelSpec):
+                set_transformer_decoder(
+                    spec.decoder,
+                    model.decoder,
+                    with_encoder_attention=False,
+                )
+
+                spec.register_vocabulary(_get_vocab(task.dictionary))
+                if not args.add_bos_token:
+                    spec.config.bos_token = spec.config.eos_token
+
+            else:
+                set_transformer_encoder(spec.encoder, model.encoder)
+                set_transformer_decoder(spec.decoder, model.decoder)
+
+                spec.register_source_vocabulary(_get_vocab(task.source_dictionary))
+                spec.register_target_vocabulary(_get_vocab(task.target_dictionary))
+                if self._no_default_special_tokens:
+                    spec.config.decoder_start_token = None
+                else:
+                    spec.config.decoder_start_token = spec.config.eos_token
+                    spec.config.add_source_eos = True
+
+            return spec
+
+
+def set_transformer_encoder(spec, module):
+    set_input_layers(spec, module)
+    for layer_spec, layer in zip(spec.layer, module.layers):
+        set_transformer_encoder_layer(layer_spec, layer)
+    if module.layer_norm is not None:
+        set_layer_norm(spec.layer_norm, module.layer_norm)
+    if module.layernorm_embedding is not None:
+        set_layer_norm(spec.layernorm_embedding, module.layernorm_embedding)
+
+
+def set_transformer_decoder(spec, module, with_encoder_attention=True):
+    set_input_layers(spec, module)
+    set_linear(spec.projection, module.output_projection)
+    for layer_spec, layer in zip(spec.layer, module.layers):
+        set_transformer_decoder_layer(
+            layer_spec,
+            layer,
+            with_encoder_attention=with_encoder_attention,
+        )
+    if module.layer_norm is not None:
+        set_layer_norm(spec.layer_norm, module.layer_norm)
+    if module.layernorm_embedding is not None:
+        set_layer_norm(spec.layernorm_embedding, module.layernorm_embedding)
+    if module.project_in_dim is not None:
+        set_linear(spec.project_in, module.project_in_dim)
+    if module.project_out_dim is not None:
+        set_linear(spec.project_out, module.project_out_dim)
+
+
+def set_input_layers(spec, module):
+    set_position_encodings(spec.position_encodings, module.embed_positions)
+    set_embeddings(
+        spec.embeddings[0] if isinstance(spec.embeddings, list) else spec.embeddings,
+        module.embed_tokens,
+    )
+    spec.scale_embeddings = module.embed_scale
+
+
+def set_transformer_encoder_layer(spec, module):
+    set_ffn(spec.ffn, module)
+    set_multi_head_attention(spec.self_attention, module.self_attn, self_attention=True)
+    set_layer_norm(spec.self_attention.layer_norm, module.self_attn_layer_norm)
+
+
+def set_transformer_decoder_layer(spec, module, with_encoder_attention=True):
+    set_ffn(spec.ffn, module)
+    set_multi_head_attention(spec.self_attention, module.self_attn, self_attention=True)
+    set_layer_norm(spec.self_attention.layer_norm, module.self_attn_layer_norm)
+    if with_encoder_attention:
+        set_multi_head_attention(spec.attention, module.encoder_attn)
+        set_layer_norm(spec.attention.layer_norm, module.encoder_attn_layer_norm)
+
+
+def set_ffn(spec, module):
+    set_layer_norm(spec.layer_norm, module.final_layer_norm)
+    set_linear(spec.linear_0, module.fc1)
+    set_linear(spec.linear_1, module.fc2)
+
+
+def set_multi_head_attention(spec, module, self_attention=False):
+    if self_attention:
+        split_layers = [common_spec.LinearSpec() for _ in range(3)]
+        set_linear(split_layers[0], module.q_proj)
+        set_linear(split_layers[1], module.k_proj)
+        set_linear(split_layers[2], module.v_proj)
+        utils.fuse_linear(spec.linear[0], split_layers)
+    else:
+        set_linear(spec.linear[0], module.q_proj)
+        split_layers = [common_spec.LinearSpec() for _ in range(2)]
+        set_linear(split_layers[0], module.k_proj)
+        set_linear(split_layers[1], module.v_proj)
+        utils.fuse_linear(spec.linear[1], split_layers)
+    set_linear(spec.linear[-1], module.out_proj)
+
+
+def set_layer_norm(spec, module):
+    spec.gamma = module.weight.numpy()
+    spec.beta = module.bias.numpy()
+
+
+def set_linear(spec, module):
+    spec.weight = module.weight.numpy()
+    if module.bias is not None:
+        spec.bias = module.bias.numpy()
+
+
+def set_embeddings(spec, module):
+    spec.weight = module.weight.numpy()
+
+
+def set_position_encodings(spec, module):
+    import torch
+
+    weight = module.weight if isinstance(module, torch.nn.Embedding) else module.weights
+    spec.encodings = weight.numpy()[module.padding_idx + 1 :]
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument("--model_path", required=True, help="Model path.")
+    parser.add_argument(
+        "--data_dir",
+        required=True,
+        help="Data directory containing the source and target vocabularies.",
+    )
+    parser.add_argument(
+        "--user_dir",
+        help="Directory containing custom extensions.",
+    )
+    parser.add_argument(
+        "--fixed_dictionary",
+        help="Fixed dictionary for multilingual models.",
+    )
+    parser.add_argument(
+        "--source_lang",
+        help="Source language. This argument is used to find dictionary file from `data_dir`.",
+    )
+    parser.add_argument(
+        "--target_lang",
+        help="Target language. This argument is used to find dictionary file from `data_dir`.",
+    )
+    parser.add_argument(
+        "--no_default_special_tokens",
+        action="store_true",
+        help=(
+            "Require all special tokens to be provided by the user during inference, "
+            "including the decoder start token."
+        ),
+    )
+    Converter.declare_arguments(parser)
+    args = parser.parse_args()
+    converter = FairseqConverter(
+        args.model_path,
+        args.data_dir,
+        source_lang=args.source_lang,
+        target_lang=args.target_lang,
+        fixed_dictionary=args.fixed_dictionary,
+        no_default_special_tokens=args.no_default_special_tokens,
+        user_dir=args.user_dir,
+    )
+    converter.convert_from_args(args)
+
+
+if __name__ == "__main__":
+    main()

ctranslate2/converters/marian.py

@@ -0,0 +1,315 @@
+import argparse
+import re
+
+from typing import List
+
+import numpy as np
+import yaml
+
+from ctranslate2.converters import utils
+from ctranslate2.converters.converter import Converter
+from ctranslate2.specs import common_spec, transformer_spec
+
+_SUPPORTED_ACTIVATIONS = {
+    "gelu": common_spec.Activation.GELUSigmoid,
+    "relu": common_spec.Activation.RELU,
+    "swish": common_spec.Activation.SWISH,
+}
+
+_SUPPORTED_POSTPROCESS_EMB = {"", "d", "n", "nd"}
+
+
+class MarianConverter(Converter):
+    """Converts models trained with Marian."""
+
+    def __init__(self, model_path: str, vocab_paths: List[str]):
+        """Initializes the Marian converter.
+
+        Arguments:
+          model_path: Path to the Marian model (.npz file).
+          vocab_paths: Paths to the vocabularies (.yml files).
+        """
+        self._model_path = model_path
+        self._vocab_paths = vocab_paths
+
+    def _load(self):
+        model = np.load(self._model_path)
+        config = _get_model_config(model)
+        vocabs = list(map(load_vocab, self._vocab_paths))
+
+        activation = config["transformer-ffn-activation"]
+        pre_norm = "n" in config["transformer-preprocess"]
+        postprocess_emb = config["transformer-postprocess-emb"]
+
+        check = utils.ConfigurationChecker()
+        check(config["type"] == "transformer", "Option --type must be 'transformer'")
+        check(
+            config["transformer-decoder-autoreg"] == "self-attention",
+            "Option --transformer-decoder-autoreg must be 'self-attention'",
+        )
+        check(
+            not config["transformer-no-projection"],
+            "Option --transformer-no-projection is not supported",
+        )
+        check(
+            activation in _SUPPORTED_ACTIVATIONS,
+            "Option --transformer-ffn-activation %s is not supported "
+            "(supported activations are: %s)"
+            % (activation, ", ".join(_SUPPORTED_ACTIVATIONS.keys())),
+        )
+        check(
+            postprocess_emb in _SUPPORTED_POSTPROCESS_EMB,
+            "Option --transformer-postprocess-emb %s is not supported (supported values are: %s)"
+            % (postprocess_emb, ", ".join(_SUPPORTED_POSTPROCESS_EMB)),
+        )
+
+        if pre_norm:
+            check(
+                config["transformer-preprocess"] == "n"
+                and config["transformer-postprocess"] == "da"
+                and config.get("transformer-postprocess-top", "") == "n",
+                "Unsupported pre-norm Transformer architecture, expected the following "
+                "combination of options: "
+                "--transformer-preprocess n "
+                "--transformer-postprocess da "
+                "--transformer-postprocess-top n",
+            )
+        else:
+            check(
+                config["transformer-preprocess"] == ""
+                and config["transformer-postprocess"] == "dan"
+                and config.get("transformer-postprocess-top", "") == "",
+                "Unsupported post-norm Transformer architecture, excepted the following "
+                "combination of options: "
+                "--transformer-preprocess '' "
+                "--transformer-postprocess dan "
+                "--transformer-postprocess-top ''",
+            )
+
+        check.validate()
+
+        alignment_layer = config["transformer-guided-alignment-layer"]
+        alignment_layer = -1 if alignment_layer == "last" else int(alignment_layer) - 1
+        layernorm_embedding = "n" in postprocess_emb
+
+        model_spec = transformer_spec.TransformerSpec.from_config(
+            (config["enc-depth"], config["dec-depth"]),
+            config["transformer-heads"],
+            pre_norm=pre_norm,
+            activation=_SUPPORTED_ACTIVATIONS[activation],
+            alignment_layer=alignment_layer,
+            alignment_heads=1,
+            layernorm_embedding=layernorm_embedding,
+        )
+        set_transformer_spec(model_spec, model)
+        model_spec.register_source_vocabulary(vocabs[0])
+        model_spec.register_target_vocabulary(vocabs[-1])
+        model_spec.config.add_source_eos = True
+        return model_spec
+
+
+def _get_model_config(model):
+    config = model["special:model.yml"]
+    config = config[:-1].tobytes()
+    config = yaml.safe_load(config)
+    return config
+
+
+def load_vocab(path):
+    # pyyaml skips some entries so we manually parse the vocabulary file.
+    with open(path, encoding="utf-8") as vocab:
+        tokens = []
+        token = None
+        idx = None
+        for i, line in enumerate(vocab):
+            line = line.rstrip("\n\r")
+            if not line:
+                continue
+
+            if line.startswith("? "):  # Complex key mapping (key)
+                token = line[2:]
+            elif token is not None:  # Complex key mapping (value)
+                idx = line[2:]
+            else:
+                token, idx = line.rsplit(":", 1)
+
+            if token is not None:
+                if token.startswith('"') and token.endswith('"'):
+                    # Unescape characters and remove quotes.
+                    token = re.sub(r"\\([^x])", r"\1", token)
+                    token = token[1:-1]
+                    if token.startswith("\\x"):
+                        # Convert the digraph \x to the actual escaped sequence.
+                        token = chr(int(token[2:], base=16))
+                elif token.startswith("'") and token.endswith("'"):
+                    token = token[1:-1]
+                    token = token.replace("''", "'")
+
+            if idx is not None:
+                try:
+                    idx = int(idx.strip())
+                except ValueError as e:
+                    raise ValueError(
+                        "Unexpected format at line %d: '%s'" % (i + 1, line)
+                    ) from e
+
+                tokens.append((idx, token))
+
+                token = None
+                idx = None
+
+    return [token for _, token in sorted(tokens, key=lambda item: item[0])]
+
+
+def set_transformer_spec(spec, weights):
+    set_transformer_encoder(spec.encoder, weights, "encoder")
+    set_transformer_decoder(spec.decoder, weights, "decoder")
+
+
+def set_transformer_encoder(spec, weights, scope):
+    set_common_layers(spec, weights, scope)
+    for i, layer_spec in enumerate(spec.layer):
+        set_transformer_encoder_layer(layer_spec, weights, "%s_l%d" % (scope, i + 1))
+
+
+def set_transformer_decoder(spec, weights, scope):
+    spec.start_from_zero_embedding = True
+    set_common_layers(spec, weights, scope)
+    for i, layer_spec in enumerate(spec.layer):
+        set_transformer_decoder_layer(layer_spec, weights, "%s_l%d" % (scope, i + 1))
+
+    set_linear(
+        spec.projection,
+        weights,
+        "%s_ff_logit_out" % scope,
+        reuse_weight=spec.embeddings.weight,
+    )
+
+
+def set_common_layers(spec, weights, scope):
+    embeddings_specs = spec.embeddings
+    if not isinstance(embeddings_specs, list):
+        embeddings_specs = [embeddings_specs]
+
+    set_embeddings(embeddings_specs[0], weights, scope)
+    set_position_encodings(
+        spec.position_encodings, weights, dim=embeddings_specs[0].weight.shape[1]
+    )
+    if hasattr(spec, "layernorm_embedding"):
+        set_layer_norm(
+            spec.layernorm_embedding,
+            weights,
+            "%s_emb" % scope,
+            pre_norm=True,
+        )
+    if hasattr(spec, "layer_norm"):
+        set_layer_norm(spec.layer_norm, weights, "%s_top" % scope)
+
+
+def set_transformer_encoder_layer(spec, weights, scope):
+    set_ffn(spec.ffn, weights, "%s_ffn" % scope)
+    set_multi_head_attention(
+        spec.self_attention, weights, "%s_self" % scope, self_attention=True
+    )
+
+
+def set_transformer_decoder_layer(spec, weights, scope):
+    set_ffn(spec.ffn, weights, "%s_ffn" % scope)
+    set_multi_head_attention(
+        spec.self_attention, weights, "%s_self" % scope, self_attention=True
+    )
+    set_multi_head_attention(spec.attention, weights, "%s_context" % scope)
+
+
+def set_multi_head_attention(spec, weights, scope, self_attention=False):
+    split_layers = [common_spec.LinearSpec() for _ in range(3)]
+    set_linear(split_layers[0], weights, scope, "q")
+    set_linear(split_layers[1], weights, scope, "k")
+    set_linear(split_layers[2], weights, scope, "v")
+
+    if self_attention:
+        utils.fuse_linear(spec.linear[0], split_layers)
+    else:
+        spec.linear[0].weight = split_layers[0].weight
+        spec.linear[0].bias = split_layers[0].bias
+        utils.fuse_linear(spec.linear[1], split_layers[1:])
+
+    set_linear(spec.linear[-1], weights, scope, "o")
+    set_layer_norm_auto(spec.layer_norm, weights, "%s_Wo" % scope)
+
+
+def set_ffn(spec, weights, scope):
+    set_layer_norm_auto(spec.layer_norm, weights, "%s_ffn" % scope)
+    set_linear(spec.linear_0, weights, scope, "1")
+    set_linear(spec.linear_1, weights, scope, "2")
+
+
+def set_layer_norm_auto(spec, weights, scope):
+    try:
+        set_layer_norm(spec, weights, scope, pre_norm=True)
+    except KeyError:
+        set_layer_norm(spec, weights, scope)
+
+
+def set_layer_norm(spec, weights, scope, pre_norm=False):
+    suffix = "_pre" if pre_norm else ""
+    spec.gamma = weights["%s_ln_scale%s" % (scope, suffix)].squeeze()
+    spec.beta = weights["%s_ln_bias%s" % (scope, suffix)].squeeze()
+
+
+def set_linear(spec, weights, scope, suffix="", reuse_weight=None):
+    weight = weights.get("%s_W%s" % (scope, suffix))
+
+    if weight is None:
+        weight = weights.get("%s_Wt%s" % (scope, suffix), reuse_weight)
+    else:
+        weight = weight.transpose()
+
+    spec.weight = weight
+
+    bias = weights.get("%s_b%s" % (scope, suffix))
+    if bias is not None:
+        spec.bias = bias.squeeze()
+
+
+def set_embeddings(spec, weights, scope):
+    spec.weight = weights.get("%s_Wemb" % scope)
+    if spec.weight is None:
+        spec.weight = weights.get("Wemb")
+
+
+def set_position_encodings(spec, weights, dim=None):
+    spec.encodings = weights.get("Wpos", _make_sinusoidal_position_encodings(dim))
+
+
+def _make_sinusoidal_position_encodings(dim, num_positions=2048):
+    positions = np.arange(num_positions)
+    timescales = np.power(10000, 2 * (np.arange(dim) // 2) / dim)
+    position_enc = np.expand_dims(positions, 1) / np.expand_dims(timescales, 0)
+    table = np.zeros_like(position_enc)
+    table[:, : dim // 2] = np.sin(position_enc[:, 0::2])
+    table[:, dim // 2 :] = np.cos(position_enc[:, 1::2])
+    return table
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument(
+        "--model_path", required=True, help="Path to the model .npz file."
+    )
+    parser.add_argument(
+        "--vocab_paths",
+        required=True,
+        nargs="+",
+        help="List of paths to the YAML vocabularies.",
+    )
+    Converter.declare_arguments(parser)
+    args = parser.parse_args()
+    converter = MarianConverter(args.model_path, args.vocab_paths)
+    converter.convert_from_args(args)
+
+
+if __name__ == "__main__":
+    main()