utils.py

from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast
import os
import json


def update_tokenizer_vocab(
    model_path: str,
    new_vocab: dict[str, int],
) -> None:
    """Update the tokenizer's vocabulary to only include the tokens in keep_indices.

    Args:
        tokenizer (PreTrainedTokenizer): The tokenizer to update.
        keep_indices (list[int]): List of token IDs to keep in the vocabulary.
    """
    new_vocab_size = len(new_vocab)
    tokenizer_json_path = os.path.join(model_path, "tokenizer.json")
    with open(tokenizer_json_path, "r", encoding="utf-8") as f:
        tokenizer_data = json.load(f)
    if "model" in tokenizer_data and "vocab" in tokenizer_data["model"]:
        print(
            f" Replacing old vocab of size {len(tokenizer_data['model']['vocab'])}..."
        )
        tokenizer_data["model"]["vocab"] = new_vocab
        tokenizer_data["model"][
            "vocab_size"
        ] = new_vocab_size  # Update model vocab size, if present
    else:
        raise ValueError(
            "Could not find 'model' or 'vocab' key in tokenizer.json structure. Cannot proceed."
        )
    if "vocab_size" in tokenizer_data:
        tokenizer_data["vocab_size"] = new_vocab_size
    if "merges" in tokenizer_data["model"]:
        original_merges = tokenizer_data["model"]["merges"]
        print(f" - Filtering {len(original_merges)} merge rules...")
        
        valid_merges = []
        # Creates a set for faster lookups
        vocab_set = set(new_vocab.keys())
        
        for merge_rule in original_merges:
            # Merges are arrays of strings like ["t", "h"] for the merge "th"
            
            # A merge is valid ONLY if both parts exist in our new vocabulary
            if merge_rule[0] in vocab_set and merge_rule[1] in vocab_set and merge_rule[0] + merge_rule[1] in vocab_set:
                valid_merges.append(merge_rule)
        
        print(f" - Kept {len(valid_merges)} valid merge rules.")
        tokenizer_data["model"]["merges"] = valid_merges
    with open(tokenizer_json_path, 'w', encoding='utf-8') as f:
        json.dump(tokenizer_data, f, ensure_ascii=False, indent=2)