src/utils.py

import json
from pathlib import Path
from typing import Dict, Tuple

def save_merges(merges: Dict[Tuple[int, int], int], save_path: str):
    """Save merges dictionary to JSON file"""
    # Convert tuple keys to strings for JSON serialization
    serializable_merges = {f"{k[0]},{k[1]}": v for k, v in merges.items()}
    
    save_dir = Path(save_path)
    save_dir.mkdir(parents=True, exist_ok=True)
    
    with open(save_dir / "merges.json", "w", encoding="utf-8") as f:
        json.dump(serializable_merges, f, ensure_ascii=False, indent=2)

def load_merges(load_path: str) -> Dict[Tuple[int, int], int]:
    """Load merges dictionary from JSON file"""
    with open(Path(load_path) / "merges.json", "r", encoding="utf-8") as f:
        serialized_merges = json.load(f)
    
    # Convert string keys back to tuples
    merges = {tuple(map(int, k.split(","))): v for k, v in serialized_merges.items()}
    return merges