3ce187da02Raise ruff line-length to 100 and extend the local pre-commit format pipeline (def-signature magic-comma normalization, short multi-line assert collapse, kwarg '=' spacing, blank-line-after-short-import removal, adjacent string-literal / f-string+plain merge, redundant-pass pruning). Every transform re-checks the file AST and is dropped if it would differ; the whole-repo reformat is verified AST-identical per file and idempotent.
127 lines
4.3 KiB
Python
127 lines
4.3 KiB
Python
import ast
|
|
import json
|
|
import os
|
|
|
|
os.environ.setdefault("PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION", "python")
|
|
|
|
from transformers.utils import sentencepiece_model_pb2
|
|
|
|
from unsloth.tokenizer_utils import fix_sentencepiece_gguf
|
|
|
|
|
|
NORMAL, CONTROL, USER_DEFINED = 1, 3, 4
|
|
|
|
_SAVE_PY = os.path.abspath(
|
|
os.path.join(os.path.dirname(__file__), "..", "..", "unsloth", "save.py")
|
|
)
|
|
_TOK_PY = os.path.abspath(
|
|
os.path.join(os.path.dirname(__file__), "..", "..", "unsloth", "tokenizer_utils.py")
|
|
)
|
|
|
|
|
|
def _build(pieces):
|
|
m = sentencepiece_model_pb2.ModelProto()
|
|
for piece, score, typ in pieces:
|
|
p = m.pieces.add()
|
|
p.piece = piece
|
|
p.score = score
|
|
p.type = typ
|
|
return m.SerializeToString()
|
|
|
|
|
|
def _read(path):
|
|
m = sentencepiece_model_pb2.ModelProto()
|
|
with open(path, "rb") as f:
|
|
m.ParseFromString(f.read())
|
|
return [(p.piece, p.type) for p in m.pieces]
|
|
|
|
|
|
def test_user_defined_special_piece_is_not_retyped(tmp_path):
|
|
pieces = [
|
|
("<s>", 0.0, CONTROL),
|
|
("a", -1.0, NORMAL),
|
|
("<ud_special>", -1.0, USER_DEFINED),
|
|
]
|
|
(tmp_path / "tokenizer.model").write_bytes(_build(pieces))
|
|
(tmp_path / "tokenizer.json").write_text(
|
|
json.dumps({"added_tokens": [{"id": 2, "content": "<ud_special>", "special": True}]})
|
|
)
|
|
fix_sentencepiece_gguf(str(tmp_path))
|
|
got = dict(_read(str(tmp_path / "tokenizer.model")))
|
|
assert got["<ud_special>"] == USER_DEFINED
|
|
|
|
|
|
def test_malformed_entry_missing_id_does_not_raise(tmp_path):
|
|
pieces = [("<s>", 0.0, CONTROL), ("a", -1.0, NORMAL), ("<sot>", -1.0, NORMAL)]
|
|
(tmp_path / "tokenizer.model").write_bytes(_build(pieces))
|
|
(tmp_path / "tokenizer.json").write_text(
|
|
json.dumps(
|
|
{
|
|
"added_tokens": [
|
|
{"content": "no_id_entry", "special": True},
|
|
{"id": 2, "content": "<sot>", "special": True},
|
|
]
|
|
}
|
|
)
|
|
)
|
|
fix_sentencepiece_gguf(str(tmp_path))
|
|
got = dict(_read(str(tmp_path / "tokenizer.model")))
|
|
assert got["<sot>"] == CONTROL
|
|
|
|
|
|
def test_entry_with_non_int_id_is_skipped(tmp_path):
|
|
pieces = [("<s>", 0.0, CONTROL), ("a", -1.0, NORMAL)]
|
|
(tmp_path / "tokenizer.model").write_bytes(_build(pieces))
|
|
(tmp_path / "tokenizer.json").write_text(
|
|
json.dumps({"added_tokens": [{"id": "oops", "content": "x", "special": True}]})
|
|
)
|
|
before = (tmp_path / "tokenizer.model").read_bytes()
|
|
fix_sentencepiece_gguf(str(tmp_path))
|
|
after = (tmp_path / "tokenizer.model").read_bytes()
|
|
assert before == after
|
|
|
|
|
|
def test_save_py_except_clause_is_broad_exception():
|
|
with open(_SAVE_PY) as f:
|
|
tree = ast.parse(f.read())
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.FunctionDef) and node.name == "unsloth_save_pretrained_gguf":
|
|
for subnode in ast.walk(node):
|
|
if isinstance(subnode, ast.Try):
|
|
body_src = "\n".join(ast.unparse(s) for s in subnode.body)
|
|
if "fix_sentencepiece_gguf(" not in body_src:
|
|
continue
|
|
handler = subnode.handlers[0]
|
|
assert handler.type is not None
|
|
assert isinstance(handler.type, ast.Name)
|
|
assert handler.type.id == "Exception"
|
|
return
|
|
raise AssertionError(
|
|
"fix_sentencepiece_gguf try block not found in unsloth_save_pretrained_gguf"
|
|
)
|
|
|
|
|
|
def test_tokenizer_utils_uses_import_protobuf_fallback_pattern():
|
|
with open(_TOK_PY) as f:
|
|
src = f.read()
|
|
tree = ast.parse(src)
|
|
for node in ast.walk(tree):
|
|
if isinstance(node, ast.FunctionDef) and node.name == "fix_sentencepiece_gguf":
|
|
fn_src = ast.unparse(node)
|
|
assert "import_protobuf" in fn_src
|
|
return
|
|
raise AssertionError("fix_sentencepiece_gguf not found in tokenizer_utils.py")
|
|
|
|
|
|
def test_all_special_tokens_are_gated_by_tokenizer_json_not_by_type(tmp_path):
|
|
pieces = [
|
|
("<s>", 0.0, CONTROL),
|
|
("a", -1.0, NORMAL),
|
|
("<n1>", -1.0, NORMAL),
|
|
("<u1>", -1.0, USER_DEFINED),
|
|
]
|
|
(tmp_path / "tokenizer.model").write_bytes(_build(pieces))
|
|
fix_sentencepiece_gguf(str(tmp_path))
|
|
got = dict(_read(str(tmp_path / "tokenizer.model")))
|
|
assert got["<n1>"] == NORMAL
|
|
assert got["<u1>"] == USER_DEFINED
|