353 lines
12 KiB
Python
353 lines
12 KiB
Python
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()
|