ctranslate2 4.7.0__cp314-cp314-macosx_11_0_arm64.whl

ctranslate2/converters/opennmt_tf.py

@@ -0,0 +1,455 @@
+import argparse
+import copy
+import os
+
+from typing import Optional, Union
+
+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.GELUTanh,
+    "relu": common_spec.Activation.RELU,
+    "swish": common_spec.Activation.SWISH,
+}
+
+
+class OpenNMTTFConverter(Converter):
+    """Converts OpenNMT-tf models."""
+
+    @classmethod
+    def from_config(
+        cls,
+        config: Union[str, dict],
+        auto_config: bool = False,
+        checkpoint_path: Optional[str] = None,
+        model: Optional[str] = None,
+    ):
+        """Creates the converter from the configuration.
+
+        Arguments:
+          config: Path to the YAML configuration, or a dictionary with the loaded configuration.
+          auto_config: Whether the model automatic configuration values should be used.
+          checkpoint_path: Path to the checkpoint or checkpoint directory to load. If not set,
+            the latest checkpoint from the model directory is loaded.
+          model: If the model instance cannot be resolved from the model directory, this argument
+             can be set to either the name of the model in the catalog or the path to the model
+             configuration.
+
+        Returns:
+          A :class:`ctranslate2.converters.OpenNMTTFConverter` instance.
+        """
+        from opennmt import config as config_util
+        from opennmt.utils.checkpoint import Checkpoint
+
+        if isinstance(config, str):
+            config = config_util.load_config([config])
+        else:
+            config = copy.deepcopy(config)
+
+        if model is None:
+            model = config_util.load_model(config["model_dir"])
+        elif os.path.exists(model):
+            model = config_util.load_model_from_file(model)
+        else:
+            model = config_util.load_model_from_catalog(model)
+
+        if auto_config:
+            config_util.merge_config(config, model.auto_config())
+
+        data_config = config_util.try_prefix_paths(config["model_dir"], config["data"])
+        model.initialize(data_config)
+
+        checkpoint = Checkpoint.from_config(config, model)
+        checkpoint_path = checkpoint.restore(checkpoint_path=checkpoint_path)
+        if checkpoint_path is None:
+            raise RuntimeError("No checkpoint was restored")
+
+        model.create_variables()
+        return cls(model)
+
+    def __init__(self, model):
+        """Initializes the converter.
+
+        Arguments:
+          model: An initialized and fully-built ``opennmt.models.Model`` instance.
+        """
+        self._model = model
+
+    def _load(self):
+        import opennmt
+
+        if isinstance(self._model, opennmt.models.LanguageModel):
+            spec_builder = TransformerDecoderSpecBuilder()
+        else:
+            spec_builder = TransformerSpecBuilder()
+
+        return spec_builder(self._model)
+
+
+class TransformerSpecBuilder:
+    def __call__(self, model):
+        import opennmt
+
+        check = utils.ConfigurationChecker()
+        check(
+            isinstance(model, opennmt.models.Transformer),
+            "Only Transformer models are supported",
+        )
+        check.validate()
+
+        check(
+            isinstance(model.encoder, opennmt.encoders.SelfAttentionEncoder),
+            "Parallel encoders are not supported",
+        )
+        check(
+            isinstance(
+                model.features_inputter,
+                (opennmt.inputters.WordEmbedder, opennmt.inputters.ParallelInputter),
+            ),
+            "Source inputter must be a WordEmbedder or a ParallelInputter",
+        )
+        check.validate()
+
+        mha = model.encoder.layers[0].self_attention.layer
+        ffn = model.encoder.layers[0].ffn.layer
+        with_relative_position = mha.maximum_relative_position is not None
+        activation_name = ffn.inner.activation.__name__
+
+        check(
+            activation_name in _SUPPORTED_ACTIVATIONS,
+            "Activation %s is not supported (supported activations are: %s)"
+            % (activation_name, ", ".join(_SUPPORTED_ACTIVATIONS.keys())),
+        )
+        check(
+            with_relative_position != bool(model.encoder.position_encoder),
+            "Relative position representation and position encoding cannot be both enabled "
+            "or both disabled",
+        )
+        check(
+            model.decoder.attention_reduction
+            != opennmt.layers.MultiHeadAttentionReduction.AVERAGE_ALL_LAYERS,
+            "Averaging all multi-head attention matrices is not supported",
+        )
+
+        source_inputters = _get_inputters(model.features_inputter)
+        target_inputters = _get_inputters(model.labels_inputter)
+        num_source_embeddings = len(source_inputters)
+        if num_source_embeddings == 1:
+            embeddings_merge = common_spec.EmbeddingsMerge.CONCAT
+        else:
+            reducer = model.features_inputter.reducer
+            embeddings_merge = None
+            if reducer is not None:
+                if isinstance(reducer, opennmt.layers.ConcatReducer):
+                    embeddings_merge = common_spec.EmbeddingsMerge.CONCAT
+                elif isinstance(reducer, opennmt.layers.SumReducer):
+                    embeddings_merge = common_spec.EmbeddingsMerge.ADD
+
+            check(
+                all(
+                    isinstance(inputter, opennmt.inputters.WordEmbedder)
+                    for inputter in source_inputters
+                ),
+                "All source inputters must WordEmbedders",
+            )
+            check(
+                embeddings_merge is not None,
+                "Unsupported embeddings reducer %s" % reducer,
+            )
+
+        alignment_layer = -1
+        alignment_heads = 1
+        if (
+            model.decoder.attention_reduction
+            == opennmt.layers.MultiHeadAttentionReduction.AVERAGE_LAST_LAYER
+        ):
+            alignment_heads = 0
+
+        check.validate()
+
+        encoder_spec = transformer_spec.TransformerEncoderSpec(
+            len(model.encoder.layers),
+            model.encoder.layers[0].self_attention.layer.num_heads,
+            pre_norm=model.encoder.layer_norm is not None,
+            activation=_SUPPORTED_ACTIVATIONS[activation_name],
+            num_source_embeddings=num_source_embeddings,
+            embeddings_merge=embeddings_merge,
+            relative_position=with_relative_position,
+        )
+
+        decoder_spec = transformer_spec.TransformerDecoderSpec(
+            len(model.decoder.layers),
+            model.decoder.layers[0].self_attention.layer.num_heads,
+            pre_norm=model.decoder.layer_norm is not None,
+            activation=_SUPPORTED_ACTIVATIONS[activation_name],
+            relative_position=with_relative_position,
+            alignment_layer=alignment_layer,
+            alignment_heads=alignment_heads,
+        )
+
+        spec = transformer_spec.TransformerSpec(encoder_spec, decoder_spec)
+
+        spec.config.add_source_bos = bool(source_inputters[0].mark_start)
+        spec.config.add_source_eos = bool(source_inputters[0].mark_end)
+        for inputter in source_inputters:
+            spec.register_source_vocabulary(_load_vocab(inputter.vocabulary_file))
+        for inputter in target_inputters:
+            spec.register_target_vocabulary(_load_vocab(inputter.vocabulary_file))
+
+        self.set_transformer_encoder(
+            spec.encoder,
+            model.encoder,
+            model.features_inputter,
+        )
+        self.set_transformer_decoder(
+            spec.decoder,
+            model.decoder,
+            model.labels_inputter,
+        )
+
+        return spec
+
+    def set_transformer_encoder(self, spec, module, inputter):
+        for embedding_spec, inputter in zip(spec.embeddings, _get_inputters(inputter)):
+            self.set_embeddings(embedding_spec, inputter)
+        if module.position_encoder is not None:
+            self.set_position_encodings(
+                spec.position_encodings,
+                module.position_encoder,
+            )
+
+        for layer_spec, layer in zip(spec.layer, module.layers):
+            self.set_multi_head_attention(
+                layer_spec.self_attention,
+                layer.self_attention,
+                self_attention=True,
+            )
+
+            self.set_ffn(layer_spec.ffn, layer.ffn)
+
+        if module.layer_norm is not None:
+            self.set_layer_norm(spec.layer_norm, module.layer_norm)
+
+    def set_transformer_decoder(self, spec, module, inputter):
+        self.set_embeddings(spec.embeddings, inputter)
+        if module.position_encoder is not None:
+            self.set_position_encodings(
+                spec.position_encodings,
+                module.position_encoder,
+            )
+
+        for layer_spec, layer in zip(spec.layer, module.layers):
+            self.set_multi_head_attention(
+                layer_spec.self_attention,
+                layer.self_attention,
+                self_attention=True,
+            )
+
+            if layer.attention:
+                self.set_multi_head_attention(
+                    layer_spec.attention,
+                    layer.attention[0],
+                    self_attention=False,
+                )
+
+            self.set_ffn(layer_spec.ffn, layer.ffn)
+
+        if module.layer_norm is not None:
+            self.set_layer_norm(spec.layer_norm, module.layer_norm)
+
+        self.set_linear(spec.projection, module.output_layer)
+
+    def set_ffn(self, spec, module):
+        self.set_linear(spec.linear_0, module.layer.inner)
+        self.set_linear(spec.linear_1, module.layer.outer)
+        self.set_layer_norm_from_wrapper(spec.layer_norm, module)
+
+    def set_multi_head_attention(self, spec, module, self_attention=False):
+        split_layers = [common_spec.LinearSpec() for _ in range(3)]
+        self.set_linear(split_layers[0], module.layer.linear_queries)
+        self.set_linear(split_layers[1], module.layer.linear_keys)
+        self.set_linear(split_layers[2], module.layer.linear_values)
+
+        if self_attention:
+            utils.fuse_linear(spec.linear[0], split_layers)
+            if module.layer.maximum_relative_position is not None:
+                spec.relative_position_keys = (
+                    module.layer.relative_position_keys.numpy()
+                )
+                spec.relative_position_values = (
+                    module.layer.relative_position_values.numpy()
+                )
+        else:
+            utils.fuse_linear(spec.linear[0], split_layers[:1])
+            utils.fuse_linear(spec.linear[1], split_layers[1:])
+
+        self.set_linear(spec.linear[-1], module.layer.linear_output)
+        self.set_layer_norm_from_wrapper(spec.layer_norm, module)
+
+    def set_layer_norm_from_wrapper(self, spec, module):
+        self.set_layer_norm(
+            spec,
+            (
+                module.output_layer_norm
+                if module.input_layer_norm is None
+                else module.input_layer_norm
+            ),
+        )
+
+    def set_layer_norm(self, spec, module):
+        spec.gamma = module.gamma.numpy()
+        spec.beta = module.beta.numpy()
+
+    def set_linear(self, spec, module):
+        spec.weight = module.kernel.numpy()
+        if not module.transpose:
+            spec.weight = spec.weight.transpose()
+        if module.bias is not None:
+            spec.bias = module.bias.numpy()
+
+    def set_embeddings(self, spec, module):
+        spec.weight = module.embedding.numpy()
+
+    def set_position_encodings(self, spec, module):
+        import opennmt
+
+        if isinstance(module, opennmt.layers.PositionEmbedder):
+            spec.encodings = module.embedding.numpy()[1:]
+
+
+class TransformerDecoderSpecBuilder(TransformerSpecBuilder):
+    def __call__(self, model):
+        import opennmt
+
+        check = utils.ConfigurationChecker()
+        check(
+            isinstance(model.decoder, opennmt.decoders.SelfAttentionDecoder),
+            "Only self-attention decoders are supported",
+        )
+        check.validate()
+
+        mha = model.decoder.layers[0].self_attention.layer
+        ffn = model.decoder.layers[0].ffn.layer
+        activation_name = ffn.inner.activation.__name__
+
+        check(
+            activation_name in _SUPPORTED_ACTIVATIONS,
+            "Activation %s is not supported (supported activations are: %s)"
+            % (activation_name, ", ".join(_SUPPORTED_ACTIVATIONS.keys())),
+        )
+        check.validate()
+
+        spec = transformer_spec.TransformerDecoderModelSpec.from_config(
+            len(model.decoder.layers),
+            mha.num_heads,
+            pre_norm=model.decoder.layer_norm is not None,
+            activation=_SUPPORTED_ACTIVATIONS[activation_name],
+        )
+
+        spec.register_vocabulary(_load_vocab(model.features_inputter.vocabulary_file))
+        self.set_transformer_decoder(
+            spec.decoder,
+            model.decoder,
+            model.features_inputter,
+        )
+
+        return spec
+
+
+def _get_inputters(inputter):
+    import opennmt
+
+    return (
+        inputter.inputters
+        if isinstance(inputter, opennmt.inputters.MultiInputter)
+        else [inputter]
+    )
+
+
+def _load_vocab(vocab, unk_token="<unk>"):
+    import opennmt
+
+    if isinstance(vocab, opennmt.data.Vocab):
+        tokens = list(vocab.words)
+    elif isinstance(vocab, list):
+        tokens = list(vocab)
+    elif isinstance(vocab, str):
+        tokens = opennmt.data.Vocab.from_file(vocab).words
+    else:
+        raise TypeError("Invalid vocabulary type")
+
+    if unk_token not in tokens:
+        tokens.append(unk_token)
+    return tokens
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument("--config", help="Path to the YAML configuration.")
+    parser.add_argument(
+        "--auto_config",
+        action="store_true",
+        help="Use the model automatic configuration values.",
+    )
+    parser.add_argument(
+        "--model_path",
+        help=(
+            "Path to the checkpoint or checkpoint directory to load. If not set, "
+            "the latest checkpoint from the model directory is loaded."
+        ),
+    )
+    parser.add_argument(
+        "--model_type",
+        help=(
+            "If the model instance cannot be resolved from the model directory, "
+            "this argument can be set to either the name of the model in the catalog "
+            "or the path to the model configuration."
+        ),
+    )
+    parser.add_argument(
+        "--src_vocab",
+        help="Path to the source vocabulary (required if no configuration is set).",
+    )
+    parser.add_argument(
+        "--tgt_vocab",
+        help="Path to the target vocabulary (required if no configuration is set).",
+    )
+    Converter.declare_arguments(parser)
+    args = parser.parse_args()
+
+    config = args.config
+    if not config:
+        if not args.model_path or not args.src_vocab or not args.tgt_vocab:
+            raise ValueError(
+                "Options --model_path, --src_vocab, --tgt_vocab are required "
+                "when a configuration is not set"
+            )
+
+        model_dir = (
+            args.model_path
+            if os.path.isdir(args.model_path)
+            else os.path.dirname(args.model_path)
+        )
+        config = {
+            "model_dir": model_dir,
+            "data": {
+                "source_vocabulary": args.src_vocab,
+                "target_vocabulary": args.tgt_vocab,
+            },
+        }
+
+    converter = OpenNMTTFConverter.from_config(
+        config,
+        auto_config=args.auto_config,
+        checkpoint_path=args.model_path,
+        model=args.model_type,
+    )
+    converter.convert_from_args(args)
+
+
+if __name__ == "__main__":
+    main()

ctranslate2/converters/opus_mt.py

@@ -0,0 +1,44 @@
+import argparse
+import os
+
+import yaml
+
+from ctranslate2.converters.marian import MarianConverter
+
+
+class OpusMTConverter(MarianConverter):
+    """Converts models trained with OPUS-MT."""
+
+    def __init__(self, model_dir: str):
+        """Initializes the OPUS-MT converter.
+
+        Arguments:
+          model_dir: Path the OPUS-MT model directory.
+        """
+        with open(
+            os.path.join(model_dir, "decoder.yml"), encoding="utf-8"
+        ) as decoder_file:
+            decoder_config = yaml.safe_load(decoder_file)
+
+        model_path = os.path.join(model_dir, decoder_config["models"][0])
+        vocab_paths = [
+            os.path.join(model_dir, path) for path in decoder_config["vocabs"]
+        ]
+        super().__init__(model_path, vocab_paths)
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+    parser.add_argument(
+        "--model_dir", required=True, help="Path to the OPUS-MT model directory."
+    )
+    OpusMTConverter.declare_arguments(parser)
+    args = parser.parse_args()
+    converter = OpusMTConverter(args.model_dir)
+    converter.convert_from_args(args)
+
+
+if __name__ == "__main__":
+    main()