"""Reconstruct a HuggingFace PegasusForConditionalGeneration checkpoint from a CTranslate2 `model.bin` that has SEPARATE source and target vocabularies. This is the OpenNMT-py 3.x family (Proxecto Nos). It differs from the shared sentencepiece OpenNMT-tf models handled by ct2_to_pegasus.py: * `encoder/embeddings_0/weight` and `decoder/embeddings/weight` are different tables of different sizes. Pegasus has one vocabulary, so the two tables are concatenated as [target | source]; encoder input ids are offset by the size of the target vocabulary and the source half is made unreachable at decode time with `final_logits_bias = -1e9`; * position encodings ARE stored in the binary (`decoder/position_encodings/encodings`), so no formula has to be guessed -- but Pegasus still refuses to save the table, so `_keys_to_ignore_on_save` must be cleared; * attention and feed-forward projections carry no bias (`add_qkvbias=False`), so zeros are written where HuggingFace insists on one. Usage: python ct2_to_pegasus_dual.py """ import json import os import sys import numpy as np import torch from transformers import PegasusConfig, PegasusForConditionalGeneration from ct2_reader import read_ct2_model, dequantize from ct2_to_pegasus import ACTIVATIONS, scalar NEG = -1e9 MAX_POS = 1024 def inspect_spec(v): def nlayers(prefix): return len({int(k.split("/")[1][6:]) for k in v if k.startswith(prefix + "/layer_") and k.split("/")[1][6:].isdigit()}) tgt_size, d_model = v["decoder/embeddings/weight"].shape src_size = v["encoder/embeddings_0/weight"].shape[0] return dict( encoder_layers=nlayers("encoder"), decoder_layers=nlayers("decoder"), source_vocab_size=int(src_size), target_vocab_size=int(tgt_size), d_model=int(d_model), heads=int(scalar(v, ("num_heads", "encoder/num_heads"))), ffn_dim=int(v["encoder/layer_0/ffn/linear_0/weight"].shape[0]), pre_norm=bool(scalar(v, ("pre_norm", "encoder/pre_norm"))), activation=ACTIVATIONS[int(scalar(v, ("activation", "encoder/activation")))], layernorm_embedding=bool(scalar(v, ("layernorm_embedding", "encoder/layernorm_embedding"), False)), relative_position=any(k.endswith("relative_attention_bias") for k in v), scale_embeddings=bool(scalar(v, ("decoder/scale_embeddings",))), output_bias="decoder/projection/bias" in v, stored_positions="decoder/position_encodings/encodings" in v, attention_bias="encoder/layer_0/self_attention/linear_0/bias" in v, ) def build_state_dict(v, spec): g = lambda n: torch.from_numpy(dequantize(v, n).copy()) sd = {} d = spec["d_model"] n_tgt = spec["target_vocab_size"] zero = lambda n: torch.zeros(n) tgt_emb = g("decoder/embeddings/weight") src_emb = g("encoder/embeddings_0/weight") shared = torch.cat([tgt_emb, src_emb], dim=0) sd["model.shared.weight"] = shared sd["model.encoder.embed_tokens.weight"] = shared.clone() sd["model.decoder.embed_tokens.weight"] = shared.clone() # decoder/projection/weight is aliased to decoder/embeddings/weight sd["lm_head.weight"] = shared.clone() bias = torch.full((1, shared.shape[0]), NEG) bias[0, :n_tgt] = g("decoder/projection/bias") if spec["output_bias"] else 0.0 sd["final_logits_bias"] = bias pos = g("decoder/position_encodings/encodings")[:MAX_POS] sd["model.encoder.embed_positions.weight"] = pos.clone() sd["model.decoder.embed_positions.weight"] = pos.clone() def ln(dst, src): sd[dst + ".weight"] = g(src + "/gamma") sd[dst + ".bias"] = g(src + "/beta") def lin(dst, src, weight=None, bias=None, n=None): w = g(src + "/weight") if weight is None else weight sd[dst + ".weight"] = w if bias is None: bias = g(src + "/bias") if src and (src + "/bias") in v else zero(w.shape[0]) sd[dst + ".bias"] = bias ln("model.encoder.layer_norm", "encoder/layer_norm") ln("model.decoder.layer_norm", "decoder/layer_norm") def split_bias(src, lo, hi): return g(src + "/bias")[lo:hi] if (src + "/bias") in v else zero(hi - lo) def self_attn(dst, src): w = g(src + "/linear_0/weight") for i, part in enumerate(("q_proj", "k_proj", "v_proj")): lo, hi = i * d, (i + 1) * d lin(dst + "." + part, None, w[lo:hi], split_bias(src + "/linear_0", lo, hi)) lin(dst + ".out_proj", src + "/linear_1") for i in range(spec["encoder_layers"]): s, p = "encoder/layer_%d" % i, "model.encoder.layers.%d" % i self_attn(p + ".self_attn", s + "/self_attention") ln(p + ".self_attn_layer_norm", s + "/self_attention/layer_norm") lin(p + ".fc1", s + "/ffn/linear_0") lin(p + ".fc2", s + "/ffn/linear_1") ln(p + ".final_layer_norm", s + "/ffn/layer_norm") for i in range(spec["decoder_layers"]): s, p = "decoder/layer_%d" % i, "model.decoder.layers.%d" % i self_attn(p + ".self_attn", s + "/self_attention") ln(p + ".self_attn_layer_norm", s + "/self_attention/layer_norm") a = s + "/attention" lin(p + ".encoder_attn.q_proj", a + "/linear_0") kv = g(a + "/linear_1/weight") lin(p + ".encoder_attn.k_proj", None, kv[:d], split_bias(a + "/linear_1", 0, d)) lin(p + ".encoder_attn.v_proj", None, kv[d:], split_bias(a + "/linear_1", d, 2 * d)) lin(p + ".encoder_attn.out_proj", a + "/linear_2") ln(p + ".encoder_attn_layer_norm", a + "/layer_norm") lin(p + ".fc1", s + "/ffn/linear_0") lin(p + ".fc2", s + "/ffn/linear_1") ln(p + ".final_layer_norm", s + "/ffn/layer_norm") return sd def convert(ct2_dir, out_dir, src_vocab, tgt_vocab, ct2_cfg=None): os.makedirs(out_dir, exist_ok=True) model = read_ct2_model(os.path.join(ct2_dir, "model.bin")) v = model["variables"] spec = inspect_spec(v) spec["ct2_spec"] = "%s rev %d, binary_version %d" % ( model["spec"], model["revision"], model["binary_version"]) cfg_json = ct2_cfg or {} spec["source_eos"] = bool(cfg_json.get("add_source_eos", False)) spec["source_bos"] = bool(cfg_json.get("add_source_bos", False)) spec["decoder_start_token"] = cfg_json.get("decoder_start_token", "") print("CT2 spec:", json.dumps(spec, indent=2)) assert spec["pre_norm"] and not spec["relative_position"] assert not spec["layernorm_embedding"] assert spec["stored_positions"] assert len(src_vocab) == spec["source_vocab_size"] assert len(tgt_vocab) == spec["target_vocab_size"] n_tgt = spec["target_vocab_size"] cfg = PegasusConfig( vocab_size=n_tgt + spec["source_vocab_size"], d_model=spec["d_model"], encoder_layers=spec["encoder_layers"], decoder_layers=spec["decoder_layers"], encoder_attention_heads=spec["heads"], decoder_attention_heads=spec["heads"], encoder_ffn_dim=spec["ffn_dim"], decoder_ffn_dim=spec["ffn_dim"], max_position_embeddings=MAX_POS, activation_function=spec["activation"], scale_embedding=spec["scale_embeddings"], dropout=0.0, attention_dropout=0.0, activation_dropout=0.0, pad_token_id=tgt_vocab.index(""), bos_token_id=tgt_vocab.index(""), eos_token_id=tgt_vocab.index(""), decoder_start_token_id=tgt_vocab.index(spec["decoder_start_token"]), forced_eos_token_id=None, max_length=512, num_beams=4, tie_word_embeddings=True, static_position_embeddings=True, ) hf = PegasusForConditionalGeneration(cfg) sd = build_state_dict(v, spec) missing, unexpected = hf.load_state_dict(sd, strict=False) print("missing:", missing, "unexpected:", unexpected) assert not missing and not unexpected hf.eval() for cls in (PegasusForConditionalGeneration,): cls._keys_to_ignore_on_save = None hf.save_pretrained(out_dir, safe_serialization=True) check = PegasusForConditionalGeneration.from_pretrained(out_dir) err = (check.model.encoder.embed_positions.weight - hf.model.encoder.embed_positions.weight).abs().max().item() print("position table round-trip max abs err:", err) assert err < 1e-6, "position table was not persisted" json.dump({"source_vocab": src_vocab, "target_vocab": tgt_vocab, "source_offset": n_tgt, "target_vocab_size": n_tgt}, open(os.path.join(out_dir, "nos_vocab.json"), "w"), ensure_ascii=False) return spec if __name__ == "__main__": d = sys.argv[1] sv = json.load(open(os.path.join(d, "source_vocabulary.json"), encoding="utf-8")) tv = json.load(open(os.path.join(d, "target_vocabulary.json"), encoding="utf-8")) convert(d, sys.argv[2], sv, tv)