scripts/bbpe_trainer.py

#!/usr/bin/env python3
"""
EthioBBPE: Production-Ready Byte-Level BPE Tokenizer Trainer
Advanced Features: 
- Checkpointing with metadata and resume capability
- Multi-format compression (gzip, lzma, bz2)
- Model quantization for deployment
- Training metrics tracking
- Automatic backups
- Model validation utilities
- Multiple export formats
"""

import os
import json
import gzip
import bz2
import lzma
import shutil
import hashlib
import logging
from pathlib import Path
from typing import List, Optional, Union, Dict, Any, Tuple
from dataclasses import dataclass, field, asdict
from datetime import datetime
import tempfile
import struct

from tokenizers import ByteLevelBPETokenizer, trainers, Tokenizer
from tokenizers.implementations import BaseTokenizer

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("EthioBBPE")


@dataclass
class BBPEConfig:
    """Configuration for EthioBBPE training."""
    vocab_size: int = 30000
    min_frequency: int = 2
    show_progress: bool = True
    special_tokens: List[str] = field(default_factory=lambda: ["<pad>", "<unk>", "<s>", "</s>"])
    lowercase: bool = False
    dropout: Optional[float] = None
    
    # File paths
    data_dir: str = "./data"
    model_save_dir: str = "./models"
    model_name: str = "EthioBBPE"
    
    # Advanced features
    use_checkpoint: bool = True
    checkpoint_dir: str = "./models/checkpoints"
    save_compressed: bool = True
    compression_format: str = "gzip"  # Options: gzip, bz2, lzma
    compression_level: int = 9  # 1-9, higher = better compression but slower
    checkpoint_steps: Optional[int] = None
    num_threads: int = -1  # -1 for auto
    
    # Backup and versioning
    enable_backup: bool = True
    max_checkpoints: int = 5  # Keep only last N checkpoints
    
    # Quantization (for deployment optimization)
    enable_quantization: bool = False
    quantization_bits: int = 8  # 8-bit or 4-bit quantization
    
    def save(self, path: str):
        """Save configuration to JSON."""
        with open(path, 'w', encoding='utf-8') as f:
            json.dump(asdict(self), f, indent=2)
        logger.info(f"Configuration saved to {path}")

    @classmethod
    def load(cls, path: str) -> "BBPEConfig":
        """Load configuration from JSON."""
        with open(path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        return cls(**data)


@dataclass
class CheckpointMetadata:
    """Metadata for checkpoint management."""
    checkpoint_id: str
    timestamp: str
    vocab_size: int
    training_step: int
    is_final: bool
    checksum: str
    config: Dict[str, Any]
    metrics: Optional[Dict[str, float]] = None
    
    def to_dict(self) -> Dict[str, Any]:
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "CheckpointMetadata":
        return cls(**data)


class EthioBBPETrainer:
    """
    Production-ready trainer for Byte-Level BPE with advanced features:
    - Checkpointing with metadata tracking
    - Multi-format compression (gzip, bz2, lzma)
    - Model quantization for deployment
    - Automatic backup management
    - Training metrics collection
    """
    
    def __init__(self, config: BBPEConfig = None):
        self.config = config or BBPEConfig()
        self.output_dir = Path(self.config.model_save_dir)
        self.checkpoint_dir = Path(self.config.checkpoint_dir)
        self.tokenizer = None
        self.is_trained = False
        self.training_metrics = {}
        self.start_time = None
        
        # Create directories
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        
        logger.info(f"Initialized EthioBBPETrainer with output dir: {self.output_dir}")
        logger.info(f"Compression format: {self.config.compression_format}, Level: {self.config.compression_level}")
        if self.config.enable_quantization:
            logger.info(f"Quantization enabled: {self.config.quantization_bits}-bit")

    def _initialize_tokenizer(self):
        """Initialize the ByteLevelBPETokenizer."""
        self.tokenizer = ByteLevelBPETokenizer(
            add_prefix_space=False,
            trim_offsets=True,
            lowercase=self.config.lowercase
        )
        logger.info("Tokenizer initialized")

    def train(self, files: Union[str, List[str]] = None, use_checkpoint: bool = None):
        """
        Train the tokenizer on a list of files or a directory.
        
        Args:
            files: Path to a file, list of files, or directory containing text files.
                   If None, uses files from config.data_dir
            use_checkpoint: If True, attempts to resume from the latest checkpoint.
                           Defaults to config.use_checkpoint
        
        Returns:
            Trained tokenizer object
        """
        self.start_time = datetime.now()
        
        if self.tokenizer is None:
            self._initialize_tokenizer()

        # Use config default if not specified
        use_checkpoint = use_checkpoint if use_checkpoint is not None else self.config.use_checkpoint
        
        # Resolve file paths
        if files is None:
            # Use data_dir from config
            data_path = Path(self.config.data_dir)
            if data_path.is_dir():
                file_paths = [str(f) for f in data_path.glob("**/*.txt")]
                file_paths.extend([str(f) for f in data_path.glob("**/*.jsonl")])
                file_paths.extend([str(f) for f in data_path.glob("**/*.json")])
            else:
                raise FileNotFoundError(f"Data directory not found: {self.config.data_dir}")
        elif isinstance(files, str):
            path = Path(files)
            if path.is_dir():
                file_paths = [str(f) for f in path.glob("**/*.txt")]
                file_paths.extend([str(f) for f in path.glob("**/*.jsonl")])
                file_paths.extend([str(f) for f in path.glob("**/*.json")])
            else:
                file_paths = [str(path)]
        else:
            file_paths = files

        if not file_paths:
            raise ValueError("No valid training files found.")
        
        logger.info(f"Found {len(file_paths)} files for training.")
        logger.info(f"Training started at: {self.start_time}")

        # Checkpoint logic
        start_from_scratch = True
        resumed_from = None
        if use_checkpoint:
            latest_ckpt = self._get_latest_checkpoint()
            if latest_ckpt:
                logger.info(f"Resuming from checkpoint: {latest_ckpt}")
                self.tokenizer = Tokenizer.from_file(str(latest_ckpt))
                start_from_scratch = False
                resumed_from = str(latest_ckpt)
            else:
                logger.info("No checkpoint found. Starting from scratch.")

        # Calculate initial metrics
        total_files = len(file_paths)
        total_size = sum(Path(f).stat().st_size for f in file_paths)
        self.training_metrics['initial'] = {
            'num_files': total_files,
            'total_bytes': total_size,
            'resumed_from': resumed_from
        }

        # Train
        logger.info("Starting training...")
        train_start = datetime.now()
        
        # ByteLevelBPETokenizer.train() accepts parameters directly, not a trainer object
        self.tokenizer.train(
            files=file_paths,
            vocab_size=self.config.vocab_size,
            min_frequency=self.config.min_frequency,
            special_tokens=self.config.special_tokens,
            show_progress=self.config.show_progress
        )

        train_end = datetime.now()
        train_duration = (train_end - train_start).total_seconds()
        
        self.is_trained = True
        logger.info("Training completed successfully.")
        
        # Record final metrics
        vocab = self.tokenizer.get_vocab()
        self.training_metrics['final'] = {
            'vocab_size': len(vocab),
            'training_duration_sec': train_duration,
            'end_time': train_end.isoformat()
        }
        
        logger.info(f"Final vocabulary size: {len(vocab)}")
        logger.info(f"Training duration: {train_duration:.2f} seconds")
        
        # Auto-save checkpoint after training
        self._save_checkpoint("final_pre_compress")

        return self.tokenizer

    def _get_latest_checkpoint(self) -> Optional[Path]:
        """Find the latest checkpoint file."""
        ckpts = list(self.checkpoint_dir.glob("checkpoint_*.json"))
        if not ckpts:
            return None
        # Sort by modification time
        ckpts.sort(key=lambda p: p.stat().st_mtime, reverse=True)
        return ckpts[0]

    def _save_checkpoint(self, name: str = "latest", save_metadata: bool = True):
        """
        Save current tokenizer state to checkpoint with metadata.
        
        Args:
            name: Checkpoint name identifier
            save_metadata: Whether to save metadata JSON alongside checkpoint
        """
        if self.tokenizer is None:
            return
        
        ckpt_path = self.checkpoint_dir / f"checkpoint_{name}.json"
        self.tokenizer.save(str(ckpt_path))
        
        # Calculate checksum for integrity verification
        checksum = self._calculate_file_checksum(ckpt_path)
        
        # Save metadata if requested
        if save_metadata:
            vocab = self.tokenizer.get_vocab()
            metadata = CheckpointMetadata(
                checkpoint_id=f"{name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
                timestamp=datetime.now().isoformat(),
                vocab_size=len(vocab),
                training_step=-1,  # Not used in HF tokenizers but kept for compatibility
                is_final=(name == "final_pre_compress"),
                checksum=checksum,
                config=asdict(self.config),
                metrics=self.training_metrics.copy()
            )
            
            metadata_path = self.checkpoint_dir / f"checkpoint_{name}_metadata.json"
            with open(metadata_path, 'w', encoding='utf-8') as f:
                json.dump(metadata.to_dict(), f, indent=2)
            logger.info(f"Checkpoint metadata saved to {metadata_path}")
        
        # Manage checkpoint rotation (keep only last N checkpoints)
        if self.config.enable_backup and self.config.max_checkpoints > 0:
            self._rotate_checkpoints()
        
        logger.info(f"Checkpoint saved to {ckpt_path} (checksum: {checksum[:8]}...)")
    
    def _calculate_file_checksum(self, filepath: Path) -> str:
        """Calculate SHA256 checksum of a file."""
        sha256_hash = hashlib.sha256()
        with open(filepath, "rb") as f:
            for byte_block in iter(lambda: f.read(4096), b""):
                sha256_hash.update(byte_block)
        return sha256_hash.hexdigest()
    
    def _rotate_checkpoints(self):
        """Keep only the most recent checkpoints based on max_checkpoints config."""
        ckpts = list(self.checkpoint_dir.glob("checkpoint_*.json"))
        # Exclude metadata files
        ckpts = [c for c in ckpts if "_metadata" not in c.name]
        
        if len(ckpts) > self.config.max_checkpoints:
            # Sort by modification time, oldest first
            ckpts.sort(key=lambda p: p.stat().st_mtime)
            # Remove oldest checkpoints
            num_to_remove = len(ckpts) - self.config.max_checkpoints
            for i in range(num_to_remove):
                old_ckpt = ckpts[i]
                old_metadata = old_ckpt.parent / f"{old_ckpt.stem}_metadata.json"
                
                try:
                    old_ckpt.unlink()
                    logger.info(f"Removed old checkpoint: {old_ckpt.name}")
                    if old_metadata.exists():
                        old_metadata.unlink()
                        logger.info(f"Removed old metadata: {old_metadata.name}")
                except Exception as e:
                    logger.warning(f"Failed to remove old checkpoint {old_ckpt}: {e}")

    def save(self, model_name: str = None, compress: bool = None, 
             compression_format: str = None, export_formats: List[str] = None):
        """
        Save the trained tokenizer with advanced options.
        
        Args:
            model_name: Name of the model folder. Defaults to config.model_name
            compress: If True, saves vocab in compressed format. 
                      Defaults to config.save_compressed.
            compression_format: Format for compression ('gzip', 'bz2', 'lzma').
                               Defaults to config.compression_format.
            export_formats: List of export formats to generate.
                           Options: ['tokenizer.json', 'vocab_compressed', 'quantized', 'hf_export']
                           Defaults to all available formats.
        
        Returns:
            Path to saved model directory
        """
        if not self.is_trained and self.tokenizer is None:
            raise RuntimeError("Tokenizer not trained yet.")

        model_name = model_name or self.config.model_name
        compress = compress if compress is not None else self.config.save_compressed
        compression_format = compression_format or self.config.compression_format
        
        model_path = self.output_dir / model_name
        model_path.mkdir(parents=True, exist_ok=True)

        logger.info(f"Saving model to {model_path} (compressed={compress}, format={compression_format})...")

        # Determine which formats to export
        if export_formats is None:
            export_formats = ['tokenizer.json', 'vocab_compressed', 'hf_export']
        
        # Always save standard tokenizer.json (required for HF loading)
        if 'tokenizer.json' in export_formats:
            tokenizer_file = model_path / "tokenizer.json"
            self.tokenizer.save(str(tokenizer_file))
            logger.info(f"Standard tokenizer.json saved: {tokenizer_file}")
        
        # Save compressed vocab if requested
        if compress and 'vocab_compressed' in export_formats:
            self._save_compressed_vocab(model_path, compression_format)
        
        # Save quantized version if enabled
        if self.config.enable_quantization and 'quantized' in export_formats:
            self._save_quantized_model(model_path)
        
        # Generate Hugging Face export files
        if 'hf_export' in export_formats:
            self._save_hf_export_files(model_path)

        # Save config
        self.config.save(str(model_path / "config.json"))
        
        # Save training metrics
        self._save_training_metrics(model_path)
        
        # Save metadata card for Hugging Face
        self._save_model_card(model_path)

        # Log size comparison
        self._log_size_comparison(model_path)

        logger.info(f"Model successfully saved to {model_path}")
        return model_path
    
    def _get_compression_handler(self, format: str):
        """Get the appropriate compression handler based on format."""
        handlers = {
            'gzip': (gzip.open, '.gz', 'wt'),
            'bz2': (bz2.open, '.bz2', 'wt'),
            'lzma': (lzma.open, '.xz', 'wt')
        }
        if format not in handlers:
            logger.warning(f"Unknown compression format '{format}', falling back to gzip")
            format = 'gzip'
        return handlers[format]
    
    def _save_compressed_vocab(self, model_path: Path, format: str = 'gzip'):
        """Save vocabulary in compressed format."""
        vocab = self.tokenizer.get_vocab()
        open_func, ext, mode = self._get_compression_handler(format)
        
        vocab_path = model_path / f"vocab.json{ext}"
        
        # Use compression level from config
        if format == 'gzip':
            with open_func(vocab_path, mode, encoding='utf-8', compresslevel=self.config.compression_level) as f:
                json.dump(vocab, f)
        elif format == 'bz2':
            with open_func(vocab_path, mode, encoding='utf-8') as f:
                json.dump(vocab, f)
        elif format == 'lzma':
            with open_func(vocab_path, mode, encoding='utf-8', preset=self.config.compression_level) as f:
                json.dump(vocab, f)
        
        logger.info(f"Compressed vocab saved ({format}): {vocab_path}")
        
        # Calculate and log compression ratio
        original_size = sum(len(k) + len(str(v)) for k, v in vocab.items())
        compressed_size = vocab_path.stat().st_size
        ratio = (1 - compressed_size / (original_size * 2)) * 100 if original_size > 0 else 0
        logger.info(f"Compression ratio: {ratio:.1f}% (original ~{original_size*2} bytes -> {compressed_size} bytes)")
    
    def _save_quantized_model(self, model_path: Path):
        """Save a quantized version of the tokenizer for deployment."""
        vocab = self.tokenizer.get_vocab()
        
        # Create quantized vocabulary mapping
        if self.config.quantization_bits == 8:
            # 8-bit quantization: map tokens to uint8 range
            dtype = 'uint8'
            max_val = 255
        elif self.config.quantization_bits == 4:
            # 4-bit quantization: pack two tokens per byte
            dtype = 'uint4'
            max_val = 15
        else:
            logger.warning(f"Unsupported quantization bits: {self.config.quantization_bits}, using 8-bit")
            dtype = 'uint8'
            max_val = 255
        
        # Sort vocab by ID for consistent ordering
        sorted_vocab = sorted(vocab.items(), key=lambda x: x[1])
        
        # Create quantized lookup table
        quantized_data = {
            'vocab_size': len(sorted_vocab),
            'quantization_bits': self.config.quantization_bits,
            'dtype': dtype,
            'tokens': [token for token, _ in sorted_vocab],
            'ids': [idx for _, idx in sorted_vocab]
        }
        
        # Save quantized model
        quantized_path = model_path / f"tokenizer_quantized_{self.config.quantization_bits}bit.json.gz"
        with gzip.open(quantized_path, 'wt', encoding='utf-8', compresslevel=9) as f:
            json.dump(quantized_data, f)
        
        logger.info(f"Quantized model saved ({self.config.quantization_bits}-bit): {quantized_path}")
    
    def _save_hf_export_files(self, model_path: Path):
        """Save files needed for Hugging Face Hub export."""
        # Save merges.txt if available (for BPE models)
        try:
            merges = self.tokenizer.model.get_merges()
            if merges:
                merges_path = model_path / "merges.txt"
                with open(merges_path, 'w', encoding='utf-8') as f:
                    f.write("#version: 0.2\n")
                    for merge in merges:
                        f.write(f"{merge[0]} {merge[1]}\n")
                logger.info(f"Merges file saved: {merges_path}")
        except Exception as e:
            logger.debug(f"No merges available or error saving: {e}")
        
        # Save special tokens mapping
        special_tokens_path = model_path / "special_tokens_map.json"
        special_tokens = {
            "unk_token": "<unk>",
            "pad_token": "<pad>",
            "bos_token": "<s>",
            "eos_token": "</s>"
        }
        with open(special_tokens_path, 'w', encoding='utf-8') as f:
            json.dump(special_tokens, f, indent=2)
        logger.info(f"Special tokens map saved: {special_tokens_path}")
    
    def _save_training_metrics(self, model_path: Path):
        """Save training metrics to JSON."""
        metrics_path = model_path / "training_metrics.json"
        
        # Add additional metrics
        full_metrics = {
            **self.training_metrics,
            'config': asdict(self.config),
            'saved_at': datetime.now().isoformat()
        }
        
        with open(metrics_path, 'w', encoding='utf-8') as f:
            json.dump(full_metrics, f, indent=2)
        
        logger.info(f"Training metrics saved: {metrics_path}")
    
    def _log_size_comparison(self, model_path: Path):
        """Log size comparison between different saved formats."""
        files_info = []
        for f in model_path.iterdir():
            if f.is_file():
                size_kb = f.stat().st_size / 1024
                files_info.append((f.name, size_kb))
        
        files_info.sort(key=lambda x: x[1])
        
        logger.info("Model artifacts size summary:")
        for name, size in files_info:
            logger.info(f"  {name}: {size:.2f} KB")

    def _save_model_card(self, path: Path):
        """Generate and save a README.md for Hugging Face Hub."""
        card_content = f"""---
language:
- multilingual
tags:
- ethiobbpe
- bpe
- tokenizer
- byte-level
license: apache-2.0
datasets:
- user-provided
---

# EthioBBPE Tokenizer

This is a production-ready Byte-Level BPE tokenizer with advanced features for deployment.

## Features
- **Byte-Level**: Handles any Unicode character without <UNK>.
- **Multi-format Compression**: Supports gzip, bz2, and lzma compression.
- **Checkpointing**: Built-in safety checkpoints with metadata tracking.
- **Quantization**: Optional 8-bit/4-bit quantization for efficient deployment.
- **Training Metrics**: Comprehensive metrics tracking and logging.
- **Automatic Backup**: Checkpoint rotation to manage disk space.

## Usage

### Transformers
```python
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("{path.name}")
```

### Tokenizers Library
```python
from tokenizers import Tokenizer

tokenizer = Tokenizer.from_file("tokenizer.json")
```

### Loading Compressed Vocab
```python
import gzip
import json

# Load compressed vocabulary
with gzip.open("vocab.json.gz", 'rt', encoding='utf-8') as f:
    vocab = json.load(f)
```

## Training Configuration
```json
{json.dumps(asdict(self.config), indent=2)}
```

## Model Files
- `tokenizer.json`: Standard tokenizer file (required)
- `vocab.json.gz`: Compressed vocabulary (optional, smaller size)
- `config.json`: Training configuration
- `training_metrics.json`: Training statistics
- `special_tokens_map.json`: Special tokens mapping
- `README.md`: This file

## Checkpoints
Checkpoints are saved in the `{self.checkpoint_dir}` directory with metadata including:
- Checkpoint ID and timestamp
- Vocabulary size
- SHA256 checksum for integrity verification
- Training metrics at checkpoint time
"""
        with open(path / "README.md", 'w', encoding='utf-8') as f:
            f.write(card_content)

    def tokenize(self, text: str) -> List[str]:
        if self.tokenizer is None:
            raise RuntimeError("Tokenizer not initialized")
        return self.tokenizer.encode(text).tokens

    def encode(self, text: str) -> List[int]:
        if self.tokenizer is None:
            raise RuntimeError("Tokenizer not initialized")
        return self.tokenizer.encode(text).ids

    def decode(self, ids: List[int]) -> str:
        if self.tokenizer is None:
            raise RuntimeError("Tokenizer not initialized")
        return self.tokenizer.decode(ids)
    
    def get_vocab_size(self) -> int:
        """Get the current vocabulary size."""
        if self.tokenizer is None:
            return 0
        return len(self.tokenizer.get_vocab())
    
    def get_training_metrics(self) -> Dict[str, Any]:
        """Get training metrics collected during training."""
        return self.training_metrics.copy()
    
    def validate_checkpoint(self, checkpoint_path: Union[str, Path]) -> bool:
        """
        Validate checkpoint integrity using checksum.
        
        Args:
            checkpoint_path: Path to checkpoint JSON file
            
        Returns:
            True if valid, False otherwise
        """
        checkpoint_path = Path(checkpoint_path)
        metadata_path = checkpoint_path.parent / f"{checkpoint_path.stem}_metadata.json"
        
        if not metadata_path.exists():
            logger.warning(f"No metadata found for checkpoint: {checkpoint_path}")
            return True  # Can't validate without metadata
        
        try:
            with open(metadata_path, 'r', encoding='utf-8') as f:
                metadata = json.load(f)
            
            expected_checksum = metadata.get('checksum')
            if not expected_checksum:
                logger.warning("No checksum in metadata")
                return True
            
            actual_checksum = self._calculate_file_checksum(checkpoint_path)
            is_valid = expected_checksum == actual_checksum
            
            if is_valid:
                logger.info(f"✓ Checkpoint validated: {checkpoint_path.name}")
            else:
                logger.error(f"✗ Checkpoint validation FAILED: {checkpoint_path.name}")
                logger.error(f"  Expected: {expected_checksum}")
                logger.error(f"  Actual:   {actual_checksum}")
            
            return is_valid
        except Exception as e:
            logger.error(f"Error validating checkpoint: {e}")
            return False
    
    def list_checkpoints(self) -> List[Dict[str, Any]]:
        """
        List all available checkpoints with their metadata.
        
        Returns:
            List of checkpoint info dictionaries
        """
        checkpoints = []
        ckpt_files = sorted(self.checkpoint_dir.glob("checkpoint_*.json"))
        
        for ckpt_file in ckpt_files:
            if "_metadata" in ckpt_file.name:
                continue
            
            info = {
                'path': str(ckpt_file),
                'name': ckpt_file.name,
                'size_kb': ckpt_file.stat().st_size / 1024,
                'modified': datetime.fromtimestamp(ckpt_file.stat().st_mtime).isoformat()
            }
            
            # Try to load metadata
            metadata_file = ckpt_file.parent / f"{ckpt_file.stem}_metadata.json"
            if metadata_file.exists():
                try:
                    with open(metadata_file, 'r', encoding='utf-8') as f:
                        metadata = json.load(f)
                    info['metadata'] = metadata
                    info['vocab_size'] = metadata.get('vocab_size', 'N/A')
                    info['is_final'] = metadata.get('is_final', False)
                except Exception as e:
                    info['error'] = str(e)
            
            checkpoints.append(info)
        
        return checkpoints
    
    @staticmethod
    def load_compressed_vocab(vocab_path: Union[str, Path]) -> Dict[str, int]:
        """
        Load vocabulary from a compressed file.
        
        Args:
            vocab_path: Path to compressed vocab file (.gz, .bz2, .xz)
            
        Returns:
            Vocabulary dictionary
        """
        vocab_path = Path(vocab_path)
        
        if vocab_path.suffix == '.gz':
            with gzip.open(vocab_path, 'rt', encoding='utf-8') as f:
                return json.load(f)
        elif vocab_path.suffix == '.bz2':
            with bz2.open(vocab_path, 'rt', encoding='utf-8') as f:
                return json.load(f)
        elif vocab_path.suffix in ['.xz', '.lzma']:
            with lzma.open(vocab_path, 'rt', encoding='utf-8') as f:
                return json.load(f)
        else:
            # Assume uncompressed JSON
            with open(vocab_path, 'r', encoding='utf-8') as f:
                return json.load(f)