96 lines
2.9 KiB
Python
96 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
PROVENANCE_SCHEMA_VERSION = 4
|
|
|
|
_SHARD_SCHEMA = {
|
|
"support": "teacher_topk_plus_other_bucket",
|
|
"layout": "chunked_sample_lists",
|
|
"logprobs_dtype": "float16",
|
|
"ids_dtype": "int32",
|
|
"other_logprob_dtype": "float16",
|
|
}
|
|
|
|
_SPECIAL_TOKEN_ID_FIELDS = (
|
|
"bos_token_id",
|
|
"eos_token_id",
|
|
"pad_token_id",
|
|
"unk_token_id",
|
|
"cls_token_id",
|
|
"sep_token_id",
|
|
"mask_token_id",
|
|
"additional_special_tokens_ids",
|
|
)
|
|
|
|
def normalize_config_revision(value: str | None) -> str | None:
|
|
if value is None:
|
|
return None
|
|
stripped = value.strip()
|
|
return stripped or None
|
|
|
|
def canonical_revision(value: str | None) -> str:
|
|
return normalize_config_revision(value) or "unversioned"
|
|
|
|
def sha256_file(path: str | Path, chunk_size: int = 1 << 20) -> str:
|
|
digest = hashlib.sha256()
|
|
with open(path, "rb") as handle:
|
|
while True:
|
|
chunk = handle.read(chunk_size)
|
|
if not chunk:
|
|
break
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
def _special_token_ids(tokenizer) -> dict[str, Any]:
|
|
snapshot: dict[str, Any] = {}
|
|
for field in _SPECIAL_TOKEN_ID_FIELDS:
|
|
value = getattr(tokenizer, field, None)
|
|
if isinstance(value, tuple):
|
|
value = list(value)
|
|
snapshot[field] = value
|
|
return snapshot
|
|
|
|
def build_tokenizer_contract(tokenizer) -> dict[str, Any]:
|
|
canonical = {
|
|
"tokenizer_class": tokenizer.__class__.__name__,
|
|
"full_vocab_size": len(tokenizer),
|
|
"special_token_ids": _special_token_ids(tokenizer),
|
|
"vocab": dict(sorted(tokenizer.get_vocab().items())),
|
|
}
|
|
encoded = json.dumps(
|
|
canonical,
|
|
sort_keys=True,
|
|
separators=(",", ":"),
|
|
ensure_ascii=True,
|
|
).encode("utf-8")
|
|
return {
|
|
"tokenizer_class": canonical["tokenizer_class"],
|
|
"full_vocab_size": canonical["full_vocab_size"],
|
|
"special_token_ids": canonical["special_token_ids"],
|
|
"fingerprint": hashlib.sha256(encoded).hexdigest(),
|
|
}
|
|
|
|
def build_shard_schema() -> dict[str, str]:
|
|
return dict(_SHARD_SCHEMA)
|
|
|
|
def collect_model_vocab_sizes(model) -> dict[str, int]:
|
|
sizes: dict[str, int] = {}
|
|
|
|
config_size = getattr(getattr(model, "config", None), "vocab_size", None)
|
|
if isinstance(config_size, int):
|
|
sizes["config"] = config_size
|
|
|
|
input_embeddings = model.get_input_embeddings()
|
|
if input_embeddings is not None and getattr(input_embeddings, "weight", None) is not None:
|
|
sizes["input_embeddings"] = int(input_embeddings.weight.shape[0])
|
|
|
|
output_embeddings = model.get_output_embeddings()
|
|
if output_embeddings is not None and getattr(output_embeddings, "weight", None) is not None:
|
|
sizes["output_embeddings"] = int(output_embeddings.weight.shape[0])
|
|
|
|
return sizes
|