Spaces:
Running on Zero
Running on Zero
File size: 5,371 Bytes
414b4fe | 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 | """Create a PaDoc-ready checkpoint from a standard image-text model."""
from __future__ import annotations
import argparse
import json
import logging
from pathlib import Path
import torch
from transformers import AutoModelForImageTextToText, AutoProcessor
from .constants import (
DEFAULT_FORK_TOKEN_MAP,
DEFAULT_SPECIAL_TOKENS,
PADOC_CONFIG_KEY,
PADOC_FORK_MAP_KEY,
PADOC_SPECIAL_TOKENS_KEY,
)
logger = logging.getLogger(__name__)
def get_padoc_metadata(model_or_config) -> dict:
config = getattr(model_or_config, "config", model_or_config)
metadata = getattr(config, PADOC_CONFIG_KEY, None)
if metadata is None and hasattr(config, "text_config"):
metadata = getattr(config.text_config, PADOC_CONFIG_KEY, None)
if not isinstance(metadata, dict) or not metadata.get(PADOC_FORK_MAP_KEY):
raise ValueError("Checkpoint has no padoc.fork_token_map metadata.")
return metadata
def get_fork_token_map(model_or_config) -> dict[str, str]:
return dict(get_padoc_metadata(model_or_config)[PADOC_FORK_MAP_KEY])
def _initialize_new_rows(model, token_ids: list[int], old_vocab_size: int, seed: int) -> None:
if not token_ids:
return
with torch.no_grad(), torch.random.fork_rng():
torch.manual_seed(seed)
input_weights = model.get_input_embeddings().weight
old_input = input_weights[:old_vocab_size].float()
input_mean = old_input.mean(0)
input_std = old_input.std(0)
for token_id in token_ids:
row = input_mean + torch.randn_like(input_mean) * input_std
input_weights[token_id].copy_(row.to(input_weights.dtype))
output = model.get_output_embeddings()
if output is not None and output.weight is not input_weights:
output_weights = output.weight
old_output = output_weights[:old_vocab_size].float()
output_mean = old_output.mean(0)
output_std = old_output.std(0)
for token_id in token_ids:
row = output_mean + torch.randn_like(output_mean) * output_std
output_weights[token_id].copy_(row.to(output_weights.dtype))
def preprocess_model(
base_model: str | Path,
output_dir: str | Path,
*,
special_tokens: list[str] | None = None,
fork_token_map: dict[str, str] | None = None,
dtype: torch.dtype = torch.bfloat16,
seed: int = 42,
) -> Path:
"""Register atomic fork tokens and persist their mapping in config.json."""
special_tokens = list(special_tokens or DEFAULT_SPECIAL_TOKENS)
fork_token_map = dict(fork_token_map or DEFAULT_FORK_TOKEN_MAP)
referenced = set(fork_token_map) | set(fork_token_map.values())
if not referenced <= set(special_tokens):
missing = sorted(referenced - set(special_tokens))
raise ValueError(f"Fork map references tokens absent from special_tokens: {missing}")
model = AutoModelForImageTextToText.from_pretrained(str(base_model), dtype=dtype)
processor = AutoProcessor.from_pretrained(str(base_model))
tokenizer = processor.tokenizer
old_vocab_size = len(tokenizer)
new_tokens = [
token
for token in special_tokens
if len(tokenizer.encode(token, add_special_tokens=False)) != 1
]
if new_tokens:
tokenizer.add_special_tokens({"additional_special_tokens": new_tokens})
model.resize_token_embeddings(len(tokenizer))
new_ids = [tokenizer.encode(token, add_special_tokens=False)[0] for token in new_tokens]
_initialize_new_rows(model, new_ids, old_vocab_size, seed)
for token in special_tokens:
ids = tokenizer.encode(token, add_special_tokens=False)
if len(ids) != 1:
raise ValueError(f"Special token {token!r} is not atomic: {ids}")
metadata = {
PADOC_SPECIAL_TOKENS_KEY: special_tokens,
PADOC_FORK_MAP_KEY: fork_token_map,
}
setattr(model.config, PADOC_CONFIG_KEY, metadata)
if hasattr(model.config, "text_config"):
setattr(model.config.text_config, PADOC_CONFIG_KEY, metadata)
output_path = Path(output_dir).expanduser().resolve()
output_path.mkdir(parents=True, exist_ok=True)
model.save_pretrained(output_path)
processor.save_pretrained(output_path)
config_path = output_path / "config.json"
with config_path.open(encoding="utf-8") as handle:
config = json.load(handle)
config[PADOC_CONFIG_KEY] = metadata
with config_path.open("w", encoding="utf-8") as handle:
json.dump(config, handle, indent=2, ensure_ascii=False)
handle.write("\n")
logger.info("Saved PaDoc-ready checkpoint to %s", output_path)
return output_path
def main(argv: list[str] | None = None) -> None:
parser = argparse.ArgumentParser(description="Create a PaDoc-ready checkpoint.")
parser.add_argument("--base-model", required=True)
parser.add_argument("--output", required=True)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--dtype", choices=("bfloat16", "float32"), default="bfloat16")
args = parser.parse_args(argv)
logging.basicConfig(level=logging.INFO)
preprocess_model(
args.base_model,
args.output,
seed=args.seed,
dtype=torch.bfloat16 if args.dtype == "bfloat16" else torch.float32,
)
if __name__ == "__main__":
main()
|