| |
| """MorphBPE JSON v1 serialization.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| from typing import Any |
|
|
| SCHEMA_VERSION = 1 |
|
|
|
|
| def dump(tokenizer: Any, path: str, extra_config: dict[str, Any] | None = None) -> None: |
| """Write a MorphBPETokenizer to disk using v1 schema.""" |
| merges_sorted = sorted(tokenizer.merges.items(), key=lambda item: item[1][1]) |
| merges_arr = [ |
| {"pair": [left, right], "merged": merged, "rank": rank} |
| for (left, right), (merged, rank) in merges_sorted |
| ] |
| payload = { |
| "version": SCHEMA_VERSION, |
| "type": "morphbpe", |
| "vocab": tokenizer.vocab, |
| "merges": merges_arr, |
| "special_tokens": {"unk": "<unk>"}, |
| "config": { |
| "boundary_constrained": True, |
| "min_boundary_confidence": getattr( |
| tokenizer, "min_boundary_confidence", 0.60 |
| ), |
| "seed_alphabet": getattr(tokenizer, "seed_alphabet", False), |
| **(extra_config or {}), |
| }, |
| } |
| with open(path, "w", encoding="utf-8") as f: |
| json.dump(payload, f, ensure_ascii=False, indent=2) |
|
|
|
|
| def load(path: str) -> dict[str, Any]: |
| with open(path, "r", encoding="utf-8") as f: |
| payload = json.load(f) |
| version = payload.get("version") |
| if version is not None and int(version) != SCHEMA_VERSION: |
| raise ValueError( |
| f"unsupported morphbpe schema version: {version}; expected {SCHEMA_VERSION}" |
| ) |
| return payload |
|
|
|
|
| def merges_from_payload( |
| payload: dict[str, Any], |
| ) -> dict[tuple[str, str], tuple[str, int]]: |
| merges: dict[tuple[str, str], tuple[str, int]] = {} |
| for i, item in enumerate(payload.get("merges", [])): |
| if isinstance(item, dict): |
| left, right = item["pair"] |
| merged = item.get("merged", left + right) |
| rank = int(item.get("rank", i)) |
| else: |
| left, right, merged, rank = item |
| merges[(left, right)] = (merged, rank) |
| return merges |
|
|
