File size: 6,136 Bytes
aa7758f | 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 | #!/usr/bin/env python3
"""Create a portable Q-Prefer adapter directory from the validated D2 run."""
from __future__ import annotations
import argparse
import hashlib
import json
import shutil
from pathlib import Path
from transformers import AutoTokenizer
from qprefer_reward.constants import (
BASE_MODEL_ID,
BASE_MODEL_REVISION,
EXPECTED_SPECIAL_TOKEN_IDS,
PUBLISHED_ADAPTER_SHA256,
SPECIAL_TOKENS,
)
def sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--source",
type=Path,
required=True,
help="Validated D2 run or checkpoint directory",
)
parser.add_argument("--output", type=Path, required=True, help="Portable output directory")
parser.add_argument("--base-model", default=BASE_MODEL_ID)
parser.add_argument(
"--base-model-source",
default=None,
help=(
"Optional local snapshot used for offline export; the public "
"--base-model is still written to metadata"
),
)
parser.add_argument("--base-revision", default=BASE_MODEL_REVISION)
parser.add_argument(
"--model-name",
default="Q-Prefer-D2-broad-near",
help="Model name written to artifact_manifest.json",
)
parser.add_argument(
"--special-embeddings-source",
type=Path,
help=(
"Exact special_token_embeddings.safetensors to use when the source "
"is the historical D2 checkpoint, which did not save this file"
),
)
parser.add_argument("--allow-unknown-checkpoint", action="store_true")
parser.add_argument("--overwrite", action="store_true")
return parser.parse_args()
def ensure_output(path: Path, overwrite: bool) -> None:
if path.exists() and any(path.iterdir()) and not overwrite:
raise FileExistsError(f"{path} is not empty; pass --overwrite to replace release files")
path.mkdir(parents=True, exist_ok=True)
def prepare_tokenizer(args: argparse.Namespace, output: Path) -> tuple[object, tuple[int, ...]]:
model_source = args.base_model_source or args.base_model
revision_kwargs = {} if Path(model_source).exists() else {"revision": args.base_revision}
tokenizer = AutoTokenizer.from_pretrained(model_source, use_fast=False, **revision_kwargs)
tokenizer.add_special_tokens({"additional_special_tokens": list(SPECIAL_TOKENS)})
token_ids = tuple(tokenizer.convert_tokens_to_ids(list(SPECIAL_TOKENS)))
if token_ids != EXPECTED_SPECIAL_TOKEN_IDS:
raise RuntimeError(
"unexpected special token ids: "
f"expected={EXPECTED_SPECIAL_TOKEN_IDS}, observed={token_ids}"
)
tokenizer.save_pretrained(output)
return tokenizer, token_ids
def main() -> None:
args = parse_args()
source = args.source.expanduser().resolve()
output = args.output.expanduser().resolve()
adapter_weights = source / "adapter_model.safetensors"
adapter_config = source / "adapter_config.json"
if not adapter_weights.is_file() or not adapter_config.is_file():
raise FileNotFoundError(
"source must contain adapter_model.safetensors and adapter_config.json"
)
observed_sha = sha256(adapter_weights)
if observed_sha != PUBLISHED_ADAPTER_SHA256 and not args.allow_unknown_checkpoint:
raise RuntimeError(
"source is not the published D2 adapter: "
f"expected sha256={PUBLISHED_ADAPTER_SHA256}, observed={observed_sha}. "
"Pass --allow-unknown-checkpoint only for an intentional new model."
)
ensure_output(output, args.overwrite)
shutil.copy2(adapter_weights, output / adapter_weights.name)
model_card = Path(__file__).resolve().parents[1] / "MODEL_CARD.md"
if model_card.is_file():
shutil.copy2(model_card, output / "README.md")
config = json.loads(adapter_config.read_text())
config["base_model_name_or_path"] = args.base_model
(output / "adapter_config.json").write_text(json.dumps(config, indent=2) + "\n")
_, token_ids = prepare_tokenizer(args, output)
source_embeddings = source / "special_token_embeddings.safetensors"
if not source_embeddings.is_file() and args.special_embeddings_source:
source_embeddings = args.special_embeddings_source.expanduser().resolve()
if source_embeddings.is_file():
# New training runs save the exact rows used during optimization. Never
# replace those rows with a fresh base-model resize.
shutil.copy2(source_embeddings, output / "special_token_embeddings.safetensors")
else:
raise FileNotFoundError(
"special_token_embeddings.safetensors is missing. New runs produced by this "
"repository save it automatically. For the historical D2 checkpoint, pass "
"--special-embeddings-source pointing to the exact file from the published artifact; "
"freshly resizing the base model is not an exact replacement."
)
files = {}
for path in sorted(output.iterdir()):
if path.is_file() and path.name != "artifact_manifest.json":
files[path.name] = {"bytes": path.stat().st_size, "sha256": sha256(path)}
manifest = {
"format_version": 1,
"model": args.model_name,
"base_model": args.base_model,
"base_revision": args.base_revision,
"special_tokens": list(SPECIAL_TOKENS),
"special_token_ids": list(token_ids),
"supported_dimensions": ["visual_quality", "text_alignment"],
"unsupported_dimensions": ["motion_quality"],
"files": files,
}
(output / "artifact_manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
print(f"Prepared portable Q-Prefer adapter at {output}")
print(f"adapter sha256: {observed_sha}")
if __name__ == "__main__":
main()
|