File size: 6,734 Bytes
fd448dd 59412c2 fd448dd 59412c2 14531a0 fd448dd 59412c2 fd448dd 32d5d5f fd448dd 59412c2 fd448dd 59412c2 fd448dd 14531a0 fd448dd 59412c2 fd448dd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | """Convert TaoTrain checkpoint + tokenizer assets into a standard HF package."""
import json
import shutil
from pathlib import Path
import torch
IGNORED_CHECKPOINT_SUFFIXES = (
".rotary.inv_freq",
)
def normalize_checkpoint(checkpoint):
if isinstance(checkpoint, dict):
if "model_state" in checkpoint:
return checkpoint["model_state"], checkpoint.get("config", {})
if "model_state_dict" in checkpoint:
return checkpoint["model_state_dict"], checkpoint.get("config", {})
return checkpoint, {}
def infer_special_token_paths(repo_dir):
subdir_metadata = repo_dir / "tokenizer" / "tokenizer.special_tokens.json"
root_metadata = repo_dir / "tokenizer.special_tokens.json"
subdir_model = repo_dir / "tokenizer" / "tokenizer.model"
subdir_vocab = repo_dir / "tokenizer" / "tokenizer.vocab"
if subdir_metadata.exists() and subdir_model.exists() and subdir_vocab.exists():
return subdir_model, subdir_metadata, subdir_vocab
if root_metadata.exists():
return repo_dir / "tokenizer.model", root_metadata, repo_dir / "tokenizer.vocab"
return repo_dir / "tokenizer" / "tokenizer.model", subdir_metadata, repo_dir / "tokenizer" / "tokenizer.vocab"
def sanitize_model_state(model_state):
"""Drop deterministic non-persistent buffers that should not participate in HF weight export."""
sanitized = {}
ignored = []
for key, value in model_state.items():
if key.endswith(IGNORED_CHECKPOINT_SUFFIXES):
ignored.append(key)
continue
sanitized[key] = value
return sanitized, ignored
def write_clean_tokenizer_metadata(repo_dir, special_tokens):
tokenizer_config = {
"backend": "custom",
"bos_token": "<BOS>",
"eos_token": "<EOS>",
"unk_token": "<UNK>",
"pad_token": "<PAD>",
"tokenizer_class": "TaoNetTokenizer",
"model_max_length": 1000000000000000019884624838656,
"extra_special_tokens": [
token
for token in special_tokens
if token not in {"<UNK>", "<BOS>", "<EOS>", "<PAD>"}
],
}
special_tokens_map = {
"unk_token": "<UNK>",
"bos_token": "<BOS>",
"eos_token": "<EOS>",
"pad_token": "<PAD>",
"additional_special_tokens": [
token
for token in special_tokens
if token not in {"<UNK>", "<BOS>", "<EOS>", "<PAD>"}
],
}
with open(repo_dir / "tokenizer_config.json", "w", encoding="utf-8") as handle:
json.dump(tokenizer_config, handle, indent=2)
handle.write("\n")
with open(repo_dir / "special_tokens_map.json", "w", encoding="utf-8") as handle:
json.dump(special_tokens_map, handle, indent=2)
handle.write("\n")
added_tokens_path = repo_dir / "added_tokens.json"
if added_tokens_path.exists():
added_tokens_path.unlink()
def main():
from configuration_taonet import TaoNetConfig
from modeling_taonet import TaoNetForCausalLM
from tokenization_taonet import TaoNetTokenizer
repo_dir = Path(__file__).resolve().parent
checkpoint_path = repo_dir / "checkpoints" / "sft" / "final_model.pt"
checkpoint = torch.load(checkpoint_path, map_location="cpu")
model_state, train_config = normalize_checkpoint(checkpoint)
model_state, ignored_keys = sanitize_model_state(model_state)
model_config = dict(train_config.get("model", {}))
metadata_model_path, metadata_path, vocab_path = infer_special_token_paths(repo_dir)
with open(metadata_path, "r", encoding="utf-8") as handle:
metadata = json.load(handle)
special_tokens = metadata.get("special_tokens", {})
hf_config = TaoNetConfig.from_taotrain_model_config(
model_config,
vocab_size=model_config.get("vocab_size", sum(1 for _ in open(vocab_path, "r", encoding="utf-8"))),
pad_token_id=special_tokens.get("<PAD>", 3),
bos_token_id=special_tokens.get("<BOS>", 1),
eos_token_id=special_tokens.get("<EOS>", 2),
unk_token_id=special_tokens.get("<UNK>", 0),
)
hf_config.architectures = ["TaoNetForCausalLM"]
hf_config.auto_map = {
"AutoConfig": "configuration_taonet.TaoNetConfig",
"AutoModelForCausalLM": "modeling_taonet.TaoNetForCausalLM",
}
model = TaoNetForCausalLM(hf_config)
current_state = model.model.state_dict()
missing = sorted(set(current_state) - set(model_state))
unexpected = sorted(set(model_state) - set(current_state))
if missing or unexpected:
if missing:
print("Missing keys while loading model state:")
for key in missing:
print(f" - {key}")
if unexpected:
print("Unexpected keys while loading model state:")
for key in unexpected:
print(f" - {key}")
raise ValueError("Checkpoint/model key mismatch detected. Refusing to export a partial model.")
model.model.load_state_dict(model_state, strict=True)
model.tie_weights()
exported_state = model.model.state_dict()
mismatched_tensors = []
for key, value in model_state.items():
exported_value = exported_state[key]
if value.shape != exported_value.shape:
mismatched_tensors.append((key, "shape"))
continue
if value.dtype.is_floating_point:
if not torch.equal(value, exported_value.to(dtype=value.dtype)):
mismatched_tensors.append((key, "value"))
else:
if not torch.equal(value, exported_value):
mismatched_tensors.append((key, "value"))
if mismatched_tensors:
print("Tensor mismatches detected after strict load:")
for key, mismatch_type in mismatched_tensors[:20]:
print(f" - {key} ({mismatch_type})")
raise ValueError("Checkpoint tensors do not match the HF wrapper after load.")
model.save_pretrained(repo_dir, safe_serialization=False)
tokenizer = TaoNetTokenizer(
vocab_file=str(metadata_model_path),
special_tokens_file=str(metadata_path),
)
tokenizer.save_pretrained(repo_dir)
write_clean_tokenizer_metadata(repo_dir, metadata.get("configured_special_tokens", []))
shutil.copyfile(metadata_model_path, repo_dir / "tokenizer.model")
shutil.copyfile(metadata_path, repo_dir / "tokenizer.special_tokens.json")
shutil.copyfile(vocab_path, repo_dir / "tokenizer.vocab")
if ignored_keys:
print("Ignored non-persistent checkpoint buffers:")
for key in ignored_keys:
print(f" - {key}")
print(f"Saved Hugging Face package to: {repo_dir}")
if __name__ == "__main__":
main()
|