codekingpro/portable-devtools

1import json2from typing import Any3from pathlib import Path4 5from tokenizers import AddedToken, Tokenizer6 7from fastembed.image.transform.operators import Compose8 9 10def load_special_tokens(model_dir: Path) -> dict[str, Any]:11    tokens_map_path = model_dir / "special_tokens_map.json"12    if not tokens_map_path.exists():13        raise ValueError(f"Could not find special_tokens_map.json in {model_dir}")14 15    with open(str(tokens_map_path)) as tokens_map_file:16        tokens_map = json.load(tokens_map_file)17 18    return tokens_map19 20 21def load_tokenizer(model_dir: Path) -> tuple[Tokenizer, dict[str, int]]:22    config_path = model_dir / "config.json"23    if not config_path.exists():24        raise ValueError(f"Could not find config.json in {model_dir}")25 26    tokenizer_path = model_dir / "tokenizer.json"27    if not tokenizer_path.exists():28        raise ValueError(f"Could not find tokenizer.json in {model_dir}")29 30    tokenizer_config_path = model_dir / "tokenizer_config.json"31    if not tokenizer_config_path.exists():32        raise ValueError(f"Could not find tokenizer_config.json in {model_dir}")33 34    with open(str(config_path)) as config_file:35        config = json.load(config_file)36 37    with open(str(tokenizer_config_path)) as tokenizer_config_file:38        tokenizer_config = json.load(tokenizer_config_file)39        assert "model_max_length" in tokenizer_config or "max_length" in tokenizer_config, (40            "Models without model_max_length or max_length are not supported."41        )42        if "model_max_length" not in tokenizer_config:43            max_context = tokenizer_config["max_length"]44        elif "max_length" not in tokenizer_config:45            max_context = tokenizer_config["model_max_length"]46        else:47            max_context = min(tokenizer_config["model_max_length"], tokenizer_config["max_length"])48 49    tokens_map = load_special_tokens(model_dir)50 51    tokenizer = Tokenizer.from_file(str(tokenizer_path))52    tokenizer.enable_truncation(max_length=max_context)53    if not tokenizer.padding:54        tokenizer.enable_padding(55            pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]56        )57 58    for token in tokens_map.values():59        if isinstance(token, str):60            tokenizer.add_special_tokens([token])61        elif isinstance(token, dict):62            tokenizer.add_special_tokens([AddedToken(**token)])63 64    special_token_to_id: dict[str, int] = {}65 66    for token in tokens_map.values():67        if isinstance(token, str):68            special_token_to_id[token] = tokenizer.token_to_id(token)69        elif isinstance(token, dict):70            token_str = token.get("content", "")71            special_token_to_id[token_str] = tokenizer.token_to_id(token_str)72 73    return tokenizer, special_token_to_id74 75 76def load_preprocessor(model_dir: Path) -> Compose:77    preprocessor_config_path = model_dir / "preprocessor_config.json"78    if not preprocessor_config_path.exists():79        raise ValueError(f"Could not find preprocessor_config.json in {model_dir}")80 81    with open(str(preprocessor_config_path)) as preprocessor_config_file:82        preprocessor_config = json.load(preprocessor_config_file)83        transforms = Compose.from_config(preprocessor_config)84    return transforms85