Initialisation du repository de Beta
This commit is contained in:
commit
14985f6dbb
9469 changed files with 1903273 additions and 0 deletions
|
|
@ -0,0 +1,8 @@
|
|||
from ctranslate2.converters.converter import Converter
|
||||
from ctranslate2.converters.fairseq import FairseqConverter
|
||||
from ctranslate2.converters.marian import MarianConverter
|
||||
from ctranslate2.converters.openai_gpt2 import OpenAIGPT2Converter
|
||||
from ctranslate2.converters.opennmt_py import OpenNMTPyConverter
|
||||
from ctranslate2.converters.opennmt_tf import OpenNMTTFConverter
|
||||
from ctranslate2.converters.opus_mt import OpusMTConverter
|
||||
from ctranslate2.converters.transformers import TransformersConverter
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
|
|
@ -0,0 +1,109 @@
|
|||
import abc
|
||||
import argparse
|
||||
import os
|
||||
import shutil
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from ctranslate2.specs.model_spec import ACCEPTED_MODEL_TYPES, ModelSpec
|
||||
|
||||
|
||||
class Converter(abc.ABC):
|
||||
"""Base class for model converters."""
|
||||
|
||||
@staticmethod
|
||||
def declare_arguments(parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
|
||||
"""Adds common conversion options to the command line parser.
|
||||
|
||||
Arguments:
|
||||
parser: Command line argument parser.
|
||||
"""
|
||||
parser.add_argument(
|
||||
"--output_dir", required=True, help="Output model directory."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vocab_mapping", default=None, help="Vocabulary mapping file (optional)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
default=None,
|
||||
choices=ACCEPTED_MODEL_TYPES,
|
||||
help="Weight quantization type.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--force",
|
||||
action="store_true",
|
||||
help="Force conversion even if the output directory already exists.",
|
||||
)
|
||||
return parser
|
||||
|
||||
def convert_from_args(self, args: argparse.Namespace) -> str:
|
||||
"""Helper function to call :meth:`ctranslate2.converters.Converter.convert`
|
||||
with the parsed command line options.
|
||||
|
||||
Arguments:
|
||||
args: Namespace containing parsed arguments.
|
||||
|
||||
Returns:
|
||||
Path to the output directory.
|
||||
"""
|
||||
return self.convert(
|
||||
args.output_dir,
|
||||
vmap=args.vocab_mapping,
|
||||
quantization=args.quantization,
|
||||
force=args.force,
|
||||
)
|
||||
|
||||
def convert(
|
||||
self,
|
||||
output_dir: str,
|
||||
vmap: Optional[str] = None,
|
||||
quantization: Optional[str] = None,
|
||||
force: bool = False,
|
||||
) -> str:
|
||||
"""Converts the model to the CTranslate2 format.
|
||||
|
||||
Arguments:
|
||||
output_dir: Output directory where the CTranslate2 model is saved.
|
||||
vmap: Optional path to a vocabulary mapping file that will be included
|
||||
in the converted model directory.
|
||||
quantization: Weight quantization scheme (possible values are: int8, int8_float32,
|
||||
int8_float16, int8_bfloat16, int16, float16, bfloat16, float32).
|
||||
force: Override the output directory if it already exists.
|
||||
|
||||
Returns:
|
||||
Path to the output directory.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the output directory already exists and :obj:`force`
|
||||
is not set.
|
||||
NotImplementedError: If the converter cannot convert this model to the
|
||||
CTranslate2 format.
|
||||
"""
|
||||
if os.path.exists(output_dir) and not force:
|
||||
raise RuntimeError(
|
||||
"output directory %s already exists, use --force to override"
|
||||
% output_dir
|
||||
)
|
||||
|
||||
model_spec = self._load()
|
||||
if model_spec is None:
|
||||
raise NotImplementedError(
|
||||
"This model is not supported by CTranslate2 or this converter"
|
||||
)
|
||||
if vmap is not None:
|
||||
model_spec.register_vocabulary_mapping(vmap)
|
||||
|
||||
model_spec.validate()
|
||||
model_spec.optimize(quantization=quantization)
|
||||
|
||||
# Create model directory.
|
||||
if os.path.exists(output_dir):
|
||||
shutil.rmtree(output_dir)
|
||||
os.makedirs(output_dir)
|
||||
model_spec.save(output_dir)
|
||||
return output_dir
|
||||
|
||||
@abc.abstractmethod
|
||||
def _load(self):
|
||||
raise NotImplementedError()
|
||||
|
|
@ -0,0 +1,352 @@
|
|||
import argparse
|
||||
|
||||
from eole.config.run import PredictConfig
|
||||
from eole.constants import PositionEncodingType
|
||||
from eole.inputters.inputter import vocabs_to_dict
|
||||
from eole.models.model import BaseModel
|
||||
|
||||
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.GELU,
|
||||
"fast_gelu": common_spec.Activation.GELUTanh,
|
||||
"relu": common_spec.Activation.RELU,
|
||||
"gated-silu": common_spec.Activation.SWISH,
|
||||
}
|
||||
|
||||
|
||||
def _get_model_spec_seq2seq(
|
||||
config, variables, src_vocabs, tgt_vocabs, num_source_embeddings
|
||||
):
|
||||
"""Creates a model specification from the model config."""
|
||||
with_relative_position = (
|
||||
getattr(config.embeddings, "position_encoding_type", None)
|
||||
== PositionEncodingType.Relative
|
||||
)
|
||||
with_rotary = (
|
||||
getattr(config.embeddings, "position_encoding_type", None)
|
||||
== PositionEncodingType.Rotary
|
||||
)
|
||||
if with_rotary:
|
||||
raise ValueError(
|
||||
"Rotary embeddings are not supported yet for encoder/decoder models"
|
||||
)
|
||||
with_alibi = (
|
||||
getattr(config.embeddings, "position_encoding_type", None)
|
||||
== PositionEncodingType.Alibi
|
||||
)
|
||||
if with_alibi:
|
||||
raise ValueError("Alibi is not supported yet for encoder/decoder models")
|
||||
activation_fn = getattr(config, "mlp_activation_fn", "relu")
|
||||
|
||||
# Return the first head of the last layer unless the model was trained with alignments.
|
||||
if getattr(config.decoder, "lambda_align", 0) == 0:
|
||||
alignment_layer = -1
|
||||
alignment_heads = 1
|
||||
else:
|
||||
alignment_layer = config.decoder.alignment_layer
|
||||
alignment_heads = config.decoder.alignment_heads
|
||||
|
||||
num_heads = getattr(config.decoder, "heads", 8)
|
||||
# num_kv = getattr(config.decoder, "heads_kv", 0)
|
||||
# if num_kv == num_heads or num_kv == 0:
|
||||
# num_kv = None
|
||||
# rotary_dim = 0 if with_rotary else None
|
||||
# rotary_interleave = getattr(config.rope_config, "rotary_interleave", True)
|
||||
ffn_glu = activation_fn == "gated-silu"
|
||||
sliding_window = getattr(config, "sliding_window", 0)
|
||||
if sliding_window != 0:
|
||||
raise ValueError(
|
||||
"Sliding window is not suported yet for encoder/decoder models"
|
||||
)
|
||||
|
||||
model_spec = transformer_spec.TransformerSpec.from_config(
|
||||
(config.encoder.layers, config.decoder.layers),
|
||||
num_heads,
|
||||
with_relative_position=with_relative_position,
|
||||
# alibi=with_alibi,
|
||||
activation=_SUPPORTED_ACTIVATIONS[activation_fn],
|
||||
ffn_glu=ffn_glu,
|
||||
rms_norm=config.layer_norm == "rms",
|
||||
# rotary_dim=rotary_dim,
|
||||
# rotary_interleave=rotary_interleave,
|
||||
# num_heads_kv=num_kv,
|
||||
# sliding_window=sliding_window,
|
||||
alignment_layer=alignment_layer,
|
||||
alignment_heads=alignment_heads,
|
||||
num_source_embeddings=num_source_embeddings,
|
||||
# multi_query_attention=getattr(opt, "multiquery", False),
|
||||
)
|
||||
|
||||
set_transformer_spec(model_spec, variables)
|
||||
for src_vocab in src_vocabs:
|
||||
model_spec.register_source_vocabulary(src_vocab)
|
||||
for tgt_vocab in tgt_vocabs:
|
||||
model_spec.register_target_vocabulary(tgt_vocab)
|
||||
|
||||
return model_spec
|
||||
|
||||
|
||||
def _get_model_spec_lm(
|
||||
config, variables, src_vocabs, tgt_vocabs, num_source_embeddings
|
||||
):
|
||||
"""Creates a model specification from the model config."""
|
||||
with_relative_position = (
|
||||
getattr(config.embeddings, "position_encoding_type", None)
|
||||
== PositionEncodingType.Relative
|
||||
)
|
||||
with_rotary = (
|
||||
getattr(config.embeddings, "position_encoding_type", None)
|
||||
== PositionEncodingType.Rotary
|
||||
)
|
||||
with_alibi = (
|
||||
getattr(config.embeddings, "position_encoding_type", None)
|
||||
== PositionEncodingType.Alibi
|
||||
)
|
||||
activation_fn = getattr(config, "mlp_activation_fn", "relu")
|
||||
num_heads = getattr(config.decoder, "heads", 8)
|
||||
num_kv = getattr(config.decoder, "heads_kv", 0)
|
||||
if num_kv == num_heads or num_kv == 0:
|
||||
num_kv = None
|
||||
rotary_dim = 0 if with_rotary else None
|
||||
rotary_interleave = getattr(config.rope_config, "rotary_interleave", True)
|
||||
ffn_glu = activation_fn == "gated-silu"
|
||||
sliding_window = getattr(config, "sliding_window", 0)
|
||||
|
||||
model_spec = transformer_spec.TransformerDecoderModelSpec.from_config(
|
||||
config.decoder.layers,
|
||||
num_heads,
|
||||
activation=_SUPPORTED_ACTIVATIONS[activation_fn],
|
||||
ffn_glu=ffn_glu,
|
||||
with_relative_position=with_relative_position,
|
||||
alibi=with_alibi,
|
||||
rms_norm=config.layer_norm == "rms",
|
||||
rotary_dim=rotary_dim,
|
||||
rotary_interleave=rotary_interleave,
|
||||
num_heads_kv=num_kv,
|
||||
sliding_window=sliding_window,
|
||||
# multi_query_attention=getattr(opt, "multiquery", False),
|
||||
)
|
||||
|
||||
set_transformer_decoder(
|
||||
model_spec.decoder,
|
||||
variables,
|
||||
with_encoder_attention=False,
|
||||
)
|
||||
|
||||
for tgt_vocab in tgt_vocabs:
|
||||
model_spec.register_vocabulary(tgt_vocab)
|
||||
|
||||
return model_spec
|
||||
|
||||
|
||||
def get_vocabs(vocab):
|
||||
src_vocabs = [vocab["src"]]
|
||||
tgt_vocabs = [vocab["tgt"]]
|
||||
return src_vocabs, tgt_vocabs
|
||||
|
||||
|
||||
class EoleConverter(Converter):
|
||||
"""Converts models generated by OpenNMT-py."""
|
||||
|
||||
def __init__(self, model_path: str):
|
||||
"""Initializes the OpenNMT-py converter.
|
||||
|
||||
Arguments:
|
||||
model_path: Path to the OpenNMT-py PyTorch model (.pt file).
|
||||
"""
|
||||
self._model_path = model_path
|
||||
|
||||
def _load(self):
|
||||
import torch
|
||||
|
||||
config = PredictConfig(model_path=self._model_path, src="dummy")
|
||||
|
||||
vocabs, model, model_config = BaseModel.load_test_model(config)
|
||||
vocabs_dict = vocabs_to_dict(vocabs)
|
||||
|
||||
config.model = model_config
|
||||
src_vocabs, tgt_vocabs = get_vocabs(vocabs_dict)
|
||||
|
||||
if config.model.decoder.decoder_type == "transformer_lm":
|
||||
spec = _get_model_spec_lm(
|
||||
config.model,
|
||||
model.state_dict(),
|
||||
src_vocabs,
|
||||
tgt_vocabs,
|
||||
num_source_embeddings=len(src_vocabs),
|
||||
)
|
||||
else:
|
||||
spec = _get_model_spec_seq2seq(
|
||||
config.model,
|
||||
model.state_dict(),
|
||||
src_vocabs,
|
||||
tgt_vocabs,
|
||||
num_source_embeddings=len(src_vocabs),
|
||||
)
|
||||
spec.config.decoder_start_token = vocabs["decoder_start_token"]
|
||||
|
||||
spec.config.bos_token = vocabs["specials"]["bos_token"]
|
||||
spec.config.eos_token = vocabs["specials"]["eos_token"]
|
||||
spec.config.unk_token = vocabs["specials"]["unk_token"]
|
||||
spec.config.layer_norm_epsilon = getattr(config, "norm_eps", 1e-6)
|
||||
|
||||
return spec
|
||||
|
||||
|
||||
def set_transformer_spec(spec, variables):
|
||||
set_transformer_encoder(spec.encoder, variables)
|
||||
set_transformer_decoder(spec.decoder, variables)
|
||||
|
||||
|
||||
def set_transformer_encoder(spec, variables):
|
||||
set_input_layers(spec, variables, "src_emb")
|
||||
set_layer_norm(spec.layer_norm, variables, "encoder.layer_norm")
|
||||
for i, layer in enumerate(spec.layer):
|
||||
set_transformer_encoder_layer(
|
||||
layer, variables, "encoder.transformer_layers.%d" % i
|
||||
)
|
||||
|
||||
|
||||
def set_transformer_decoder(spec, variables, with_encoder_attention=True):
|
||||
set_input_layers(spec, variables, "tgt_emb")
|
||||
set_layer_norm(spec.layer_norm, variables, "decoder.layer_norm")
|
||||
for i, layer in enumerate(spec.layer):
|
||||
set_transformer_decoder_layer(
|
||||
layer,
|
||||
variables,
|
||||
"decoder.transformer_layers.%d" % i,
|
||||
with_encoder_attention=with_encoder_attention,
|
||||
)
|
||||
|
||||
set_linear(spec.projection, variables, "generator")
|
||||
|
||||
|
||||
def set_input_layers(spec, variables, scope):
|
||||
if hasattr(spec, "position_encodings"):
|
||||
set_position_encodings(
|
||||
spec.position_encodings,
|
||||
variables,
|
||||
"%s.pe" % scope,
|
||||
)
|
||||
else:
|
||||
spec.scale_embeddings = False
|
||||
|
||||
embeddings_specs = spec.embeddings
|
||||
# encoder embeddings are stored in a list(onmt/ct2 legacy with features)
|
||||
if isinstance(embeddings_specs, list):
|
||||
embeddings_specs = embeddings_specs[0]
|
||||
set_embeddings(embeddings_specs, variables, "%s.embeddings" % scope)
|
||||
|
||||
|
||||
def set_transformer_encoder_layer(spec, variables, scope):
|
||||
set_multi_head_attention(
|
||||
spec.self_attention,
|
||||
variables,
|
||||
"%s.self_attn" % scope,
|
||||
self_attention=True,
|
||||
)
|
||||
set_layer_norm(
|
||||
spec.self_attention.layer_norm, variables, "%s.input_layernorm" % scope
|
||||
)
|
||||
set_layer_norm(
|
||||
spec.ffn.layer_norm, variables, "%s.post_attention_layernorm" % scope
|
||||
)
|
||||
set_ffn(spec.ffn, variables, "%s.mlp" % scope)
|
||||
|
||||
|
||||
def set_transformer_decoder_layer(spec, variables, scope, with_encoder_attention=True):
|
||||
set_multi_head_attention(
|
||||
spec.self_attention,
|
||||
variables,
|
||||
"%s.self_attn" % scope,
|
||||
self_attention=True,
|
||||
)
|
||||
set_layer_norm(
|
||||
spec.self_attention.layer_norm, variables, "%s.input_layernorm" % scope
|
||||
)
|
||||
if with_encoder_attention:
|
||||
set_multi_head_attention(spec.attention, variables, "%s.context_attn" % scope)
|
||||
set_layer_norm(
|
||||
spec.attention.layer_norm, variables, "%s.precontext_layernorm" % scope
|
||||
)
|
||||
set_layer_norm(
|
||||
spec.ffn.layer_norm, variables, "%s.post_attention_layernorm" % scope
|
||||
)
|
||||
set_ffn(spec.ffn, variables, "%s.mlp" % scope)
|
||||
|
||||
|
||||
def set_ffn(spec, variables, scope):
|
||||
set_linear(spec.linear_0, variables, "%s.gate_up_proj" % scope)
|
||||
set_linear(spec.linear_1, variables, "%s.down_proj" % scope)
|
||||
if hasattr(spec, "linear_0_noact"):
|
||||
set_linear(spec.linear_0_noact, variables, "%s.up_proj" % scope)
|
||||
|
||||
|
||||
def set_multi_head_attention(spec, variables, scope, self_attention=False):
|
||||
if self_attention:
|
||||
split_layers = [common_spec.LinearSpec() for _ in range(3)]
|
||||
set_linear(split_layers[0], variables, "%s.linear_query" % scope)
|
||||
set_linear(split_layers[1], variables, "%s.linear_keys" % scope)
|
||||
set_linear(split_layers[2], variables, "%s.linear_values" % scope)
|
||||
utils.fuse_linear(spec.linear[0], split_layers)
|
||||
else:
|
||||
set_linear(spec.linear[0], variables, "%s.linear_query" % scope)
|
||||
split_layers = [common_spec.LinearSpec() for _ in range(2)]
|
||||
set_linear(split_layers[0], variables, "%s.linear_keys" % scope)
|
||||
set_linear(split_layers[1], variables, "%s.linear_values" % scope)
|
||||
utils.fuse_linear(spec.linear[1], split_layers)
|
||||
set_linear(spec.linear[-1], variables, "%s.final_linear" % scope)
|
||||
if hasattr(spec, "relative_position_keys"):
|
||||
spec.relative_position_keys = _get_variable(
|
||||
variables, "%s.relative_positions_embeddings.weight" % scope
|
||||
)
|
||||
spec.relative_position_values = spec.relative_position_keys
|
||||
|
||||
|
||||
def set_layer_norm(spec, variables, scope):
|
||||
try:
|
||||
spec.gamma = _get_variable(variables, "%s.weight" % scope)
|
||||
except KeyError:
|
||||
# Compatibility with older models using a custom LayerNorm module.
|
||||
spec.gamma = _get_variable(variables, "%s.a_2" % scope)
|
||||
spec.beta = _get_variable(variables, "%s.b_2" % scope)
|
||||
try:
|
||||
spec.beta = _get_variable(variables, "%s.bias" % scope)
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
|
||||
def set_linear(spec, variables, scope):
|
||||
spec.weight = _get_variable(variables, "%s.weight" % scope)
|
||||
bias = variables.get("%s.bias" % scope)
|
||||
if bias is not None:
|
||||
spec.bias = bias
|
||||
|
||||
|
||||
def set_embeddings(spec, variables, scope):
|
||||
spec.weight = _get_variable(variables, "%s.weight" % scope)
|
||||
|
||||
|
||||
def set_position_encodings(spec, variables, scope):
|
||||
spec.encodings = _get_variable(variables, "%s.pe" % scope).squeeze()
|
||||
|
||||
|
||||
def _get_variable(variables, name):
|
||||
return variables[name]
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
||||
)
|
||||
parser.add_argument("--model_path", required=True, help="Model path.")
|
||||
Converter.declare_arguments(parser)
|
||||
args = parser.parse_args()
|
||||
EoleConverter(args.model_path).convert_from_args(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -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()
|
||||
|
|
@ -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()
|
||||
|
|
@ -0,0 +1,95 @@
|
|||
import argparse
|
||||
import json
|
||||
import os
|
||||
|
||||
from ctranslate2.converters.converter import Converter
|
||||
from ctranslate2.specs import common_spec, model_spec, transformer_spec
|
||||
|
||||
|
||||
class OpenAIGPT2Converter(Converter):
|
||||
"""Converts GPT-2 models from https://github.com/openai/gpt-2."""
|
||||
|
||||
def __init__(self, model_dir: str):
|
||||
"""Initializes the OpenAI GPT-2 converter.
|
||||
|
||||
Arguments:
|
||||
model_dir: Path to the OpenAI GPT-2 model directory.
|
||||
"""
|
||||
self._model_dir = model_dir
|
||||
|
||||
def _load(self):
|
||||
import tensorflow as tf
|
||||
|
||||
reader = tf.train.load_checkpoint(self._model_dir)
|
||||
weights = {
|
||||
name: reader.get_tensor(name)
|
||||
for name in reader.get_variable_to_shape_map().keys()
|
||||
}
|
||||
|
||||
with open(os.path.join(self._model_dir, "hparams.json")) as hparams_file:
|
||||
hparams = json.load(hparams_file)
|
||||
with open(os.path.join(self._model_dir, "encoder.json")) as vocab_file:
|
||||
vocab = json.load(vocab_file)
|
||||
vocab = [
|
||||
token
|
||||
for token, index in sorted(vocab.items(), key=lambda item: item[1])
|
||||
]
|
||||
|
||||
spec = transformer_spec.TransformerDecoderModelSpec.from_config(
|
||||
hparams["n_layer"],
|
||||
hparams["n_head"],
|
||||
pre_norm=True,
|
||||
activation=common_spec.Activation.GELUTanh,
|
||||
)
|
||||
set_decoder(spec.decoder, weights, "model")
|
||||
spec.unk_token = "<|endoftext|>"
|
||||
spec.bos_token = "<|endoftext|>"
|
||||
spec.eos_token = "<|endoftext|>"
|
||||
spec.register_vocabulary(vocab)
|
||||
return spec
|
||||
|
||||
|
||||
def set_decoder(spec, weights, scope):
|
||||
spec.embeddings.weight = weights["%s/wte" % scope]
|
||||
spec.position_encodings.encodings = weights["%s/wpe" % scope]
|
||||
spec.scale_embeddings = False
|
||||
spec.projection.weight = spec.embeddings.weight
|
||||
set_layer_norm(spec.layer_norm, weights, "%s/ln_f" % scope)
|
||||
for i, layer_spec in enumerate(spec.layer):
|
||||
set_layer(layer_spec, weights, "%s/h%d" % (scope, i))
|
||||
|
||||
|
||||
def set_layer_norm(spec, weights, scope):
|
||||
spec.gamma = weights["%s/g" % scope]
|
||||
spec.beta = weights["%s/b" % scope]
|
||||
|
||||
|
||||
def set_linear(spec, weights, scope):
|
||||
spec.weight = weights["%s/w" % scope].squeeze().transpose()
|
||||
spec.bias = weights["%s/b" % scope]
|
||||
|
||||
|
||||
def set_layer(spec, weights, scope):
|
||||
set_layer_norm(spec.self_attention.layer_norm, weights, "%s/ln_1" % scope)
|
||||
set_linear(spec.self_attention.linear[0], weights, "%s/attn/c_attn" % scope)
|
||||
set_linear(spec.self_attention.linear[1], weights, "%s/attn/c_proj" % scope)
|
||||
set_layer_norm(spec.ffn.layer_norm, weights, "%s/ln_2" % scope)
|
||||
set_linear(spec.ffn.linear_0, weights, "%s/mlp/c_fc" % scope)
|
||||
set_linear(spec.ffn.linear_1, weights, "%s/mlp/c_proj" % scope)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_dir", required=True, help="Path to the model directory."
|
||||
)
|
||||
Converter.declare_arguments(parser)
|
||||
args = parser.parse_args()
|
||||
converter = OpenAIGPT2Converter(args.model_dir)
|
||||
converter.convert_from_args(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,361 @@
|
|||
import argparse
|
||||
|
||||
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.GELU,
|
||||
"fast_gelu": common_spec.Activation.GELUTanh,
|
||||
"relu": common_spec.Activation.RELU,
|
||||
"silu": common_spec.Activation.SWISH,
|
||||
}
|
||||
|
||||
_SUPPORTED_FEATURES_MERGE = {
|
||||
"concat": common_spec.EmbeddingsMerge.CONCAT,
|
||||
"sum": common_spec.EmbeddingsMerge.ADD,
|
||||
}
|
||||
|
||||
|
||||
def check_opt(opt, num_source_embeddings):
|
||||
with_relative_position = getattr(opt, "max_relative_positions", 0) > 0
|
||||
with_rotary = getattr(opt, "max_relative_positions", 0) == -1
|
||||
with_alibi = getattr(opt, "max_relative_positions", 0) == -2
|
||||
activation_fn = getattr(opt, "pos_ffn_activation_fn", "relu")
|
||||
feat_merge = getattr(opt, "feat_merge", "concat")
|
||||
self_attn_type = getattr(opt, "self_attn_type", "scaled-dot")
|
||||
|
||||
check = utils.ConfigurationChecker()
|
||||
check(
|
||||
opt.encoder_type == opt.decoder_type
|
||||
and opt.decoder_type in {"transformer", "transformer_lm"},
|
||||
"Options --encoder_type and --decoder_type must be"
|
||||
" 'transformer' or 'transformer_lm",
|
||||
)
|
||||
check(
|
||||
self_attn_type == "scaled-dot",
|
||||
"Option --self_attn_type %s is not supported (supported values are: scaled-dot)"
|
||||
% self_attn_type,
|
||||
)
|
||||
check(
|
||||
activation_fn in _SUPPORTED_ACTIVATIONS,
|
||||
"Option --pos_ffn_activation_fn %s is not supported (supported activations are: %s)"
|
||||
% (activation_fn, ", ".join(_SUPPORTED_ACTIVATIONS.keys())),
|
||||
)
|
||||
check(
|
||||
opt.position_encoding != (with_relative_position or with_rotary or with_alibi),
|
||||
"Options --position_encoding and --max_relative_positions cannot be both enabled "
|
||||
"or both disabled",
|
||||
)
|
||||
check(
|
||||
num_source_embeddings == 1 or feat_merge in _SUPPORTED_FEATURES_MERGE,
|
||||
"Option --feat_merge %s is not supported (supported merge modes are: %s)"
|
||||
% (feat_merge, " ".join(_SUPPORTED_FEATURES_MERGE.keys())),
|
||||
)
|
||||
check.validate()
|
||||
|
||||
|
||||
def _get_model_spec_seq2seq(
|
||||
opt, variables, src_vocabs, tgt_vocabs, num_source_embeddings
|
||||
):
|
||||
"""Creates a model specification from the model options."""
|
||||
with_relative_position = getattr(opt, "max_relative_positions", 0) > 0
|
||||
activation_fn = getattr(opt, "pos_ffn_activation_fn", "relu")
|
||||
feat_merge = getattr(opt, "feat_merge", "concat")
|
||||
|
||||
# Return the first head of the last layer unless the model was trained with alignments.
|
||||
if getattr(opt, "lambda_align", 0) == 0:
|
||||
alignment_layer = -1
|
||||
alignment_heads = 1
|
||||
else:
|
||||
alignment_layer = opt.alignment_layer
|
||||
alignment_heads = opt.alignment_heads
|
||||
|
||||
num_heads = getattr(opt, "heads", 8)
|
||||
|
||||
model_spec = transformer_spec.TransformerSpec.from_config(
|
||||
(opt.enc_layers, opt.dec_layers),
|
||||
num_heads,
|
||||
with_relative_position=with_relative_position,
|
||||
activation=_SUPPORTED_ACTIVATIONS[activation_fn],
|
||||
alignment_layer=alignment_layer,
|
||||
alignment_heads=alignment_heads,
|
||||
num_source_embeddings=num_source_embeddings,
|
||||
embeddings_merge=_SUPPORTED_FEATURES_MERGE[feat_merge],
|
||||
multi_query_attention=getattr(opt, "multiquery", False),
|
||||
)
|
||||
|
||||
model_spec.config.decoder_start_token = getattr(opt, "decoder_start_token", "<s>")
|
||||
|
||||
set_transformer_spec(model_spec, variables)
|
||||
for src_vocab in src_vocabs:
|
||||
model_spec.register_source_vocabulary(src_vocab)
|
||||
for tgt_vocab in tgt_vocabs:
|
||||
model_spec.register_target_vocabulary(tgt_vocab)
|
||||
|
||||
return model_spec
|
||||
|
||||
|
||||
def _get_model_spec_lm(opt, variables, src_vocabs, tgt_vocabs, num_source_embeddings):
|
||||
"""Creates a model specification from the model options."""
|
||||
with_relative_position = getattr(opt, "max_relative_positions", 0) > 0
|
||||
with_rotary = getattr(opt, "max_relative_positions", 0) == -1
|
||||
with_alibi = getattr(opt, "max_relative_positions", 0) == -2
|
||||
activation_fn = getattr(opt, "pos_ffn_activation_fn", "relu")
|
||||
num_heads = getattr(opt, "heads", 8)
|
||||
num_kv = getattr(opt, "num_kv", 0)
|
||||
if num_kv == num_heads or num_kv == 0:
|
||||
num_kv = None
|
||||
rotary_dim = 0 if with_rotary else None
|
||||
rotary_interleave = getattr(opt, "rotary_interleave", True)
|
||||
ffn_glu = activation_fn == "silu"
|
||||
sliding_window = getattr(opt, "sliding_window", 0)
|
||||
|
||||
model_spec = transformer_spec.TransformerDecoderModelSpec.from_config(
|
||||
opt.dec_layers,
|
||||
num_heads,
|
||||
activation=_SUPPORTED_ACTIVATIONS[activation_fn],
|
||||
ffn_glu=ffn_glu,
|
||||
with_relative_position=with_relative_position,
|
||||
alibi=with_alibi,
|
||||
rms_norm=opt.layer_norm == "rms",
|
||||
rotary_dim=rotary_dim,
|
||||
rotary_interleave=rotary_interleave,
|
||||
multi_query_attention=getattr(opt, "multiquery", False),
|
||||
num_heads_kv=num_kv,
|
||||
sliding_window=sliding_window,
|
||||
)
|
||||
|
||||
model_spec.config.layer_norm_epsilon = getattr(opt, "norm_eps", 1e-6)
|
||||
|
||||
set_transformer_decoder(
|
||||
model_spec.decoder,
|
||||
variables,
|
||||
with_encoder_attention=False,
|
||||
)
|
||||
|
||||
for tgt_vocab in tgt_vocabs:
|
||||
model_spec.register_vocabulary(tgt_vocab)
|
||||
|
||||
return model_spec
|
||||
|
||||
|
||||
def get_vocabs(vocab):
|
||||
if isinstance(vocab, dict) and "src" in vocab:
|
||||
if isinstance(vocab["src"], list):
|
||||
src_vocabs = [vocab["src"]]
|
||||
tgt_vocabs = [vocab["tgt"]]
|
||||
|
||||
src_feats = vocab.get("src_feats")
|
||||
if src_feats is not None:
|
||||
src_vocabs.extend(src_feats.values())
|
||||
else:
|
||||
src_vocabs = [field[1].vocab.itos for field in vocab["src"].fields]
|
||||
tgt_vocabs = [field[1].vocab.itos for field in vocab["tgt"].fields]
|
||||
else:
|
||||
# Compatibility with older models.
|
||||
src_vocabs = [vocab[0][1].itos]
|
||||
tgt_vocabs = [vocab[1][1].itos]
|
||||
|
||||
return src_vocabs, tgt_vocabs
|
||||
|
||||
|
||||
class OpenNMTPyConverter(Converter):
|
||||
"""Converts models generated by OpenNMT-py."""
|
||||
|
||||
def __init__(self, model_path: str):
|
||||
"""Initializes the OpenNMT-py converter.
|
||||
|
||||
Arguments:
|
||||
model_path: Path to the OpenNMT-py PyTorch model (.pt file).
|
||||
"""
|
||||
self._model_path = model_path
|
||||
|
||||
def _load(self):
|
||||
import torch
|
||||
|
||||
checkpoint = torch.load(
|
||||
self._model_path, map_location="cpu", weights_only=False
|
||||
)
|
||||
|
||||
src_vocabs, tgt_vocabs = get_vocabs(checkpoint["vocab"])
|
||||
|
||||
check_opt(checkpoint["opt"], num_source_embeddings=len(src_vocabs))
|
||||
|
||||
variables = checkpoint["model"]
|
||||
variables.update(
|
||||
{
|
||||
"generator.%s" % key: value
|
||||
for key, value in checkpoint["generator"].items()
|
||||
}
|
||||
)
|
||||
|
||||
if checkpoint["opt"].decoder_type == "transformer_lm":
|
||||
return _get_model_spec_lm(
|
||||
checkpoint["opt"],
|
||||
variables,
|
||||
src_vocabs,
|
||||
tgt_vocabs,
|
||||
num_source_embeddings=len(src_vocabs),
|
||||
)
|
||||
else:
|
||||
return _get_model_spec_seq2seq(
|
||||
checkpoint["opt"],
|
||||
variables,
|
||||
src_vocabs,
|
||||
tgt_vocabs,
|
||||
num_source_embeddings=len(src_vocabs),
|
||||
)
|
||||
|
||||
|
||||
def set_transformer_spec(spec, variables):
|
||||
set_transformer_encoder(spec.encoder, variables)
|
||||
set_transformer_decoder(spec.decoder, variables)
|
||||
|
||||
|
||||
def set_transformer_encoder(spec, variables):
|
||||
set_input_layers(spec, variables, "encoder")
|
||||
set_layer_norm(spec.layer_norm, variables, "encoder.layer_norm")
|
||||
for i, layer in enumerate(spec.layer):
|
||||
set_transformer_encoder_layer(layer, variables, "encoder.transformer.%d" % i)
|
||||
|
||||
|
||||
def set_transformer_decoder(spec, variables, with_encoder_attention=True):
|
||||
set_input_layers(spec, variables, "decoder")
|
||||
set_layer_norm(spec.layer_norm, variables, "decoder.layer_norm")
|
||||
for i, layer in enumerate(spec.layer):
|
||||
set_transformer_decoder_layer(
|
||||
layer,
|
||||
variables,
|
||||
"decoder.transformer_layers.%d" % i,
|
||||
with_encoder_attention=with_encoder_attention,
|
||||
)
|
||||
|
||||
try:
|
||||
set_linear(spec.projection, variables, "generator")
|
||||
except KeyError:
|
||||
# Compatibility when the generator was a nn.Sequential module.
|
||||
set_linear(spec.projection, variables, "generator.0")
|
||||
|
||||
|
||||
def set_input_layers(spec, variables, scope):
|
||||
if hasattr(spec, "position_encodings"):
|
||||
set_position_encodings(
|
||||
spec.position_encodings,
|
||||
variables,
|
||||
"%s.embeddings.make_embedding.pe" % scope,
|
||||
)
|
||||
else:
|
||||
# See https://github.com/OpenNMT/OpenNMT-py/issues/1722
|
||||
spec.scale_embeddings = False
|
||||
|
||||
embeddings_specs = spec.embeddings
|
||||
if not isinstance(embeddings_specs, list):
|
||||
embeddings_specs = [embeddings_specs]
|
||||
|
||||
for i, embeddings_spec in enumerate(embeddings_specs):
|
||||
set_embeddings(
|
||||
embeddings_spec,
|
||||
variables,
|
||||
"%s.embeddings.make_embedding.emb_luts.%d" % (scope, i),
|
||||
)
|
||||
|
||||
|
||||
def set_transformer_encoder_layer(spec, variables, scope):
|
||||
set_ffn(spec.ffn, variables, "%s.feed_forward" % scope)
|
||||
set_multi_head_attention(
|
||||
spec.self_attention,
|
||||
variables,
|
||||
"%s.self_attn" % scope,
|
||||
self_attention=True,
|
||||
)
|
||||
set_layer_norm(spec.self_attention.layer_norm, variables, "%s.layer_norm" % scope)
|
||||
|
||||
|
||||
def set_transformer_decoder_layer(spec, variables, scope, with_encoder_attention=True):
|
||||
set_ffn(spec.ffn, variables, "%s.feed_forward" % scope)
|
||||
set_multi_head_attention(
|
||||
spec.self_attention,
|
||||
variables,
|
||||
"%s.self_attn" % scope,
|
||||
self_attention=True,
|
||||
)
|
||||
set_layer_norm(spec.self_attention.layer_norm, variables, "%s.layer_norm_1" % scope)
|
||||
if with_encoder_attention:
|
||||
set_multi_head_attention(spec.attention, variables, "%s.context_attn" % scope)
|
||||
set_layer_norm(spec.attention.layer_norm, variables, "%s.layer_norm_2" % scope)
|
||||
|
||||
|
||||
def set_ffn(spec, variables, scope):
|
||||
set_layer_norm(spec.layer_norm, variables, "%s.layer_norm" % scope)
|
||||
set_linear(spec.linear_0, variables, "%s.w_1" % scope)
|
||||
set_linear(spec.linear_1, variables, "%s.w_2" % scope)
|
||||
if hasattr(spec, "linear_0_noact"):
|
||||
set_linear(spec.linear_0_noact, variables, "%s.w_3" % scope)
|
||||
|
||||
|
||||
def set_multi_head_attention(spec, variables, scope, self_attention=False):
|
||||
if self_attention:
|
||||
split_layers = [common_spec.LinearSpec() for _ in range(3)]
|
||||
set_linear(split_layers[0], variables, "%s.linear_query" % scope)
|
||||
set_linear(split_layers[1], variables, "%s.linear_keys" % scope)
|
||||
set_linear(split_layers[2], variables, "%s.linear_values" % scope)
|
||||
utils.fuse_linear(spec.linear[0], split_layers)
|
||||
else:
|
||||
set_linear(spec.linear[0], variables, "%s.linear_query" % scope)
|
||||
split_layers = [common_spec.LinearSpec() for _ in range(2)]
|
||||
set_linear(split_layers[0], variables, "%s.linear_keys" % scope)
|
||||
set_linear(split_layers[1], variables, "%s.linear_values" % scope)
|
||||
utils.fuse_linear(spec.linear[1], split_layers)
|
||||
set_linear(spec.linear[-1], variables, "%s.final_linear" % scope)
|
||||
if hasattr(spec, "relative_position_keys"):
|
||||
spec.relative_position_keys = _get_variable(
|
||||
variables, "%s.relative_positions_embeddings.weight" % scope
|
||||
)
|
||||
spec.relative_position_values = spec.relative_position_keys
|
||||
|
||||
|
||||
def set_layer_norm(spec, variables, scope):
|
||||
try:
|
||||
spec.gamma = _get_variable(variables, "%s.weight" % scope)
|
||||
except KeyError:
|
||||
# Compatibility with older models using a custom LayerNorm module.
|
||||
spec.gamma = _get_variable(variables, "%s.a_2" % scope)
|
||||
spec.beta = _get_variable(variables, "%s.b_2" % scope)
|
||||
try:
|
||||
spec.beta = _get_variable(variables, "%s.bias" % scope)
|
||||
except KeyError:
|
||||
pass
|
||||
|
||||
|
||||
def set_linear(spec, variables, scope):
|
||||
spec.weight = _get_variable(variables, "%s.weight" % scope)
|
||||
bias = variables.get("%s.bias" % scope)
|
||||
if bias is not None:
|
||||
spec.bias = bias
|
||||
|
||||
|
||||
def set_embeddings(spec, variables, scope):
|
||||
spec.weight = _get_variable(variables, "%s.weight" % scope)
|
||||
|
||||
|
||||
def set_position_encodings(spec, variables, scope):
|
||||
spec.encodings = _get_variable(variables, "%s.pe" % scope).squeeze()
|
||||
|
||||
|
||||
def _get_variable(variables, name):
|
||||
return variables[name]
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
||||
)
|
||||
parser.add_argument("--model_path", required=True, help="Model path.")
|
||||
Converter.declare_arguments(parser)
|
||||
args = parser.parse_args()
|
||||
OpenNMTPyConverter(args.model_path).convert_from_args(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -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()
|
||||
|
|
@ -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()
|
||||
File diff suppressed because it is too large
Load diff
|
|
@ -0,0 +1,127 @@
|
|||
import numpy as np
|
||||
|
||||
|
||||
def fuse_linear(spec, layers):
|
||||
if not layers:
|
||||
raise ValueError("Cannot fuse linear layers: at least one layer is required")
|
||||
|
||||
if isinstance(layers[0].weight, np.ndarray):
|
||||
concatenate = np.concatenate
|
||||
zeros = np.zeros
|
||||
else:
|
||||
import torch
|
||||
|
||||
concatenate = torch.cat
|
||||
zeros = torch.zeros
|
||||
|
||||
spec.weight = concatenate([layer.weight for layer in layers])
|
||||
|
||||
bias_dtype = None
|
||||
for layer in layers:
|
||||
if layer.has_bias():
|
||||
bias_dtype = layer.bias.dtype
|
||||
break
|
||||
|
||||
if bias_dtype is not None:
|
||||
spec.bias = concatenate(
|
||||
[
|
||||
(
|
||||
layer.bias
|
||||
if layer.has_bias()
|
||||
else zeros([layer.weight.shape[0]], dtype=bias_dtype)
|
||||
)
|
||||
for layer in layers
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def fuse_linear_prequant(spec, layers, axis):
|
||||
if not layers:
|
||||
raise ValueError("Cannot fuse linear layers: at least one layer is required")
|
||||
params = ["weight", "weight_scale", "weight_zero"]
|
||||
if isinstance(layers[0].weight, np.ndarray):
|
||||
concatenate = np.concatenate
|
||||
else:
|
||||
import torch
|
||||
|
||||
concatenate = torch.cat
|
||||
|
||||
for param in params:
|
||||
setattr(
|
||||
spec,
|
||||
param,
|
||||
concatenate([getattr(layer, param) for layer in layers], axis=axis),
|
||||
)
|
||||
|
||||
|
||||
def permute_for_sliced_rotary(weight, num_heads, rotary_dim=None):
|
||||
"""Permutes the weight to use the sliced rotary implementation."""
|
||||
if rotary_dim is not None:
|
||||
weight = weight.reshape(num_heads, weight.shape[0] // num_heads, -1)
|
||||
|
||||
rotary_weight = weight[:, :rotary_dim]
|
||||
rotary_weight = permute_for_sliced_rotary(
|
||||
rotary_weight.reshape(num_heads * rotary_dim, -1), num_heads
|
||||
).reshape(num_heads, rotary_dim, -1)
|
||||
|
||||
weight[:, :rotary_dim] = rotary_weight
|
||||
|
||||
return weight.reshape(-1, weight.shape[-1])
|
||||
|
||||
return (
|
||||
weight.reshape(num_heads, weight.shape[0] // num_heads // 2, 2, weight.shape[1])
|
||||
.swapaxes(1, 2)
|
||||
.reshape(weight.shape[0], weight.shape[1])
|
||||
)
|
||||
|
||||
|
||||
def smooth_activation(layer_norm, linear, activation_scales):
|
||||
"""Applies the activation smoothing technique described in
|
||||
https://github.com/mit-han-lab/smoothquant.
|
||||
"""
|
||||
if not isinstance(linear.weight, np.ndarray):
|
||||
linear_weight = linear.weight.numpy()
|
||||
activation_scales = activation_scales.numpy()
|
||||
else:
|
||||
linear_weight = linear.weight
|
||||
|
||||
weight_scales = np.amax(np.absolute(linear_weight), axis=0)
|
||||
weight_scales = np.maximum(weight_scales, 1e-5)
|
||||
|
||||
activation_scales = activation_scales.astype(weight_scales.dtype)
|
||||
|
||||
scales = np.sqrt(activation_scales / weight_scales)
|
||||
scales = np.maximum(scales, 1e-5)
|
||||
|
||||
if not isinstance(linear.weight, np.ndarray):
|
||||
import torch
|
||||
|
||||
scales = torch.from_numpy(scales)
|
||||
|
||||
layer_norm.gamma /= scales
|
||||
layer_norm.beta /= scales
|
||||
|
||||
linear.weight *= scales.reshape(1, -1)
|
||||
|
||||
|
||||
def raise_unsupported(reasons):
|
||||
message = (
|
||||
"The model you are trying to convert is not supported by CTranslate2. "
|
||||
"We identified the following reasons:\n"
|
||||
)
|
||||
for reason in reasons:
|
||||
message += "\n- " + reason
|
||||
raise ValueError(message)
|
||||
|
||||
|
||||
class ConfigurationChecker:
|
||||
def __init__(self):
|
||||
self._unsupported_reasons = []
|
||||
|
||||
def __call__(self, assert_condition, error_message):
|
||||
if not assert_condition:
|
||||
self._unsupported_reasons.append(error_message)
|
||||
|
||||
def validate(self):
|
||||
if self._unsupported_reasons:
|
||||
raise_unsupported(self._unsupported_reasons)
|
||||
Loading…
Add table
Add a link
Reference in a new issue