codon-model 0.0.1__py3-none-any.whl

codon/utils/mask.py

@@ -0,0 +1,266 @@
+import torch
+
+from tokenizers  import Tokenizer
+from dataclasses import dataclass
+from enum        import Enum, auto
+
+from typing import Union
+
+
+class MaskMode(Enum):
+    '''
+    Enumeration of different masking modes for TokenMask.
+
+    Each mode defines how the sequence is masked relative to the special token(s).
+    The mask values are: 0 for masked, 1 for unmasked (kept).
+
+    Attributes:
+        FIRST_MASK_PRE: Find the first occurrence of the special token.
+            Mask tokens before and including the special token (0). Keep the rest (1).
+        FIRST_MASK_POST: Find the first occurrence of the special token.
+            Keep tokens before and including the special token (1). Mask the rest (0).
+        LAST_MASK_PRE: Find the last occurrence of the special token.
+            Mask tokens before and including the special token (0). Keep the rest (1).
+        LAST_MASK_POST: Find the last occurrence of the special token.
+            Keep tokens before and including the special token (1). Mask the rest (0).
+        ALL_MASK_FIRST: Find all occurrences.
+            The first segment (ending with the special token) is masked (0), then alternates.
+        ALL_KEEP_FIRST: Find all occurrences.
+            The first segment (ending with the special token) is kept (1), then alternates.
+    '''
+    FIRST_MASK_PRE  = auto()
+    FIRST_MASK_POST = auto()
+    LAST_MASK_PRE   = auto()
+    LAST_MASK_POST  = auto()
+    ALL_MASK_FIRST  = auto()
+    ALL_KEEP_FIRST  = auto()
+
+
+def make_padding_mask(src: torch.Tensor, pad_idx: int = 0) -> torch.Tensor:
+    '''
+    Creates a padding mask.
+
+    Args:
+        src (torch.Tensor): The source sequence tensor. Shape is [B, L_src].
+        pad_idx (int, optional): The index of the padding symbol. Defaults to 0.
+
+    Returns:
+        torch.Tensor: The padding mask. Shape is [B, 1, 1, L_src].
+        True indicates the position is not padding and should be attended to.
+    '''
+    mask = (src != pad_idx).unsqueeze(1).unsqueeze(2)
+    return mask
+
+
+def make_lookahead_mask(size: int, device: torch.device = torch.device('cpu')) -> torch.Tensor:
+    '''
+    Creates a lookahead mask (lower triangular matrix).
+
+    Args:
+        size (int): The sequence length.
+        device (torch.device, optional): The device. Defaults to cpu.
+
+    Returns:
+        torch.Tensor: The lookahead mask. Shape is [size, size].
+        True indicates allowed positions to attend to (lower triangular part).
+    '''
+    mask = torch.tril(torch.ones((size, size), device=device)).bool()
+    return mask
+
+
+def make_causal_mask(tgt: torch.Tensor, pad_idx: int = 0) -> torch.Tensor:
+    '''
+    Creates a causal mask (combining padding mask and lookahead mask).
+
+    Args:
+        tgt (torch.Tensor): The target sequence tensor. Shape is [B, L_tgt].
+        pad_idx (int, optional): The index of the padding symbol. Defaults to 0.
+
+    Returns:
+        torch.Tensor: The causal mask. Shape is [B, 1, L_tgt, L_tgt].
+    '''
+    pad_mask = make_padding_mask(tgt, pad_idx)
+    seq_len = tgt.size(1)
+    lookahead_mask = make_lookahead_mask(seq_len, device=tgt.device)
+
+    # pad_mask: [B, 1, 1, L]
+    # lookahead_mask: [L, L]
+    mask = pad_mask & lookahead_mask
+    return mask
+
+
+def make_sliding_window_mask(
+    tensor: torch.Tensor, window_size: int, pad_idx: int = 0, causal: bool = True
+) -> torch.Tensor:
+    '''
+    Creates a sliding window mask.
+
+    Args:
+        tensor (torch.Tensor): The input sequence tensor. Shape is [B, L].
+        window_size (int): The window size (one-sided).
+        pad_idx (int, optional): The index of the padding symbol. Defaults to 0.
+        causal (bool, optional): Whether it is causal (unidirectional). Defaults to True.
+            If True, position i can only attend to [i - window_size, i].
+            If False, position i can attend to [i - window_size, i + window_size].
+
+    Returns:
+        torch.Tensor: The sliding window mask. Shape is [B, 1, L, L].
+    '''
+    pad_mask = make_padding_mask(tensor, pad_idx)  # [B, 1, 1, L]
+    seq_len = tensor.size(1)
+
+    ones = torch.ones((seq_len, seq_len), device=tensor.device, dtype=torch.bool)
+
+    if causal:
+        # j <= i AND j >= i - window_size
+        window_mask = torch.tril(ones, diagonal=0) & torch.triu(
+            ones, diagonal=-window_size
+        )
+    else:
+        # j <= i + window_size AND j >= i - window_size
+        window_mask = torch.tril(ones, diagonal=window_size) & torch.triu(
+            ones, diagonal=-window_size
+        )
+
+    mask = pad_mask & window_mask
+    return mask
+
+
+@dataclass
+class MaskedContent:
+    '''
+    Result of the token masking process.
+
+    Attributes:
+        content (str): The original text content.
+        tokenized (Union[list[int], torch.Tensor]): The list of token IDs or tensor.
+        mask (Union[list[int], torch.Tensor]): The mask values (0 for masked, 1 for unmasked).
+    '''
+    content: str
+    tokenized: Union[list[int], torch.Tensor]
+    mask: Union[list[int], torch.Tensor]
+
+
+class TokenMask:
+    '''
+    Handles token masking logic based on special tokens.
+    '''
+
+    def __init__(self, tokenizer: Tokenizer) -> None:
+        '''
+        Initializes the TokenMask.
+
+        Args:
+            tokenizer (Tokenizer): The configured tokenizer instance.
+        '''
+        self.tokenizer = tokenizer
+
+    def mask(
+        self,
+        content: str,
+        special_token: Union[str, int, list[Union[str, int]]],
+        mode: MaskMode = MaskMode.FIRST_MASK_PRE,
+        tensor_mask: bool = True
+    ) -> MaskedContent:
+        '''
+        Tokenizes content and generates a mask based on the specified mode.
+
+        Args:
+            content (str): The text content to tokenize and mask.
+            special_token (Union[str, int, list[Union[str, int]]]): The special token(s) to use as a separator.
+            mode (MaskMode): The masking mode. Defaults to MaskMode.FIRST_MASK_PRE.
+            tensor_mask (bool, optional): Whether to return tensors instead of lists. Defaults to True.
+
+        Returns:
+            MaskedContent: Dataclass containing the original content, token IDs, and the generated mask.
+        '''
+        encoded = self.tokenizer.encode(content)
+        ids = encoded.ids
+
+        candidates = []
+        if isinstance(special_token, list):
+            candidates = special_token
+        else:
+            candidates = [special_token]
+
+        # Determine the separator token id used in the sequence
+        sep_id = None
+        for cand in candidates:
+            tid = None
+            if isinstance(cand, str):
+                tid = self.tokenizer.token_to_id(cand)
+            elif isinstance(cand, int):
+                tid = cand
+            
+            if tid is not None and tid in ids:
+                sep_id = tid
+                break
+        
+        mask = []
+
+        if sep_id is None:
+             if mode in [MaskMode.FIRST_MASK_POST, MaskMode.LAST_MASK_POST, MaskMode.ALL_KEEP_FIRST]:
+                 mask = [1] * len(ids)
+             else:
+                 mask = [0] * len(ids)
+        else:
+            indices = [i for i, x in enumerate(ids) if x == sep_id]
+            
+            if mode == MaskMode.FIRST_MASK_PRE:
+                # 1. Find first, mask pre (0), keep post (1).
+                # [0, 0, sep, 1, 1]
+                idx = indices[0]
+                mask = [0] * (idx + 1) + [1] * (len(ids) - idx - 1)
+
+            elif mode == MaskMode.FIRST_MASK_POST:
+                # 2. Find first, keep pre (1), mask post (0).
+                # [1, 1, sep, 0, 0]
+                idx = indices[0]
+                mask = [1] * (idx + 1) + [0] * (len(ids) - idx - 1)
+
+            elif mode == MaskMode.LAST_MASK_PRE:
+                # 3. Find last, mask pre (0), keep post (1).
+                idx = indices[-1]
+                mask = [0] * (idx + 1) + [1] * (len(ids) - idx - 1)
+
+            elif mode == MaskMode.LAST_MASK_POST:
+                # 4. Find last, keep pre (1), mask post (0).
+                idx = indices[-1]
+                mask = [1] * (idx + 1) + [0] * (len(ids) - idx - 1)
+
+            elif mode == MaskMode.ALL_MASK_FIRST:
+                # 5. All, segments. First seg mask (0), second unmask (1)...
+                # Segments end at sep.
+                current_val = 0
+                last_idx = 0
+                mask = []
+                for idx in indices:
+                    # chunk includes sep
+                    chunk_len = idx - last_idx + 1
+                    mask.extend([current_val] * chunk_len)
+                    current_val = 1 - current_val # toggle
+                    last_idx = idx + 1
+                
+                # Remaining part
+                if last_idx < len(ids):
+                    mask.extend([current_val] * (len(ids) - last_idx))
+
+            elif mode == MaskMode.ALL_KEEP_FIRST:
+                # 6. All, segments. First seg keep (1), second mask (0)...
+                current_val = 1
+                last_idx = 0
+                mask = []
+                for idx in indices:
+                    chunk_len = idx - last_idx + 1
+                    mask.extend([current_val] * chunk_len)
+                    current_val = 1 - current_val
+                    last_idx = idx + 1
+                
+                if last_idx < len(ids):
+                    mask.extend([current_val] * (len(ids) - last_idx))
+
+        if tensor_mask:
+            ids = torch.tensor(ids)
+            mask = torch.tensor(mask)
+
+        return MaskedContent(content=content, tokenized=ids, mask=mask)

codon/utils/safecode.py

@@ -0,0 +1,24 @@
+import random
+import string
+
+
+def safecode(length: int = 4, exclude_confusing: bool = False) -> str:
+    '''
+    Generates a random safe code consisting of letters and digits.
+
+    Args:
+        length (int): The length of the code to generate. Defaults to 4.
+        exclude_confusing (bool): If True, excludes confusing characters
+            ('0oO1iIlLq9g') to reduce human error. Defaults to False.
+
+    Returns:
+        str: The generated random code.
+    '''
+    characters = string.ascii_letters + string.digits
+
+    if exclude_confusing:
+        confusing_chars = '0oO1iIlLq9g'
+        characters = ''.join(c for c in characters if c not in confusing_chars)
+
+    code = ''.join(random.choices(characters, k=length))
+    return code

codon/utils/seed.py

@@ -0,0 +1,75 @@
+import torch
+import numpy as np
+import random
+import os
+
+from typing import Any, Optional
+
+def seed_everything(seed: int = 42, strict: bool = False, warn_only: bool = True) -> None:
+    '''
+    Sets all random seeds to ensure reproducibility of PyTorch experiments.
+
+    Args:
+        seed (int): The random seed value. Defaults to 42.
+        strict (bool): Whether to enable strict deterministic mode.
+            If True, calls torch.use_deterministic_algorithms(True),
+            which may raise errors for certain operations that do not support
+            deterministic algorithms and may reduce training speed.
+        warn_only (bool): If True, only warns when deterministic algorithms
+            are not available. Defaults to True.
+    '''
+    import codon
+    codon.__seed__ = seed
+    random.seed(seed)
+    np.random.seed(seed)
+    
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    
+    if torch.cuda.is_available():
+        torch.backends.cudnn.benchmark = False 
+        torch.backends.cudnn.deterministic = True 
+    if strict:
+        try:
+            torch.use_deterministic_algorithms(True, warn_only=warn_only)
+            os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
+            print(f'[Info] Strict deterministic mode enabled. (seed={seed})')
+        except AttributeError:
+            print('[Warning] torch.use_deterministic_algorithms is not available in your PyTorch version.')
+    else:
+        print(f'[Info] Random seed set as {seed}')
+
+def get_seed() -> Optional[int]:
+    '''
+    Retrieves the global random seed.
+
+    Returns:
+        Optional[int]: The current seed value, or None if not set.
+    '''
+    import codon
+    return codon.__seed__
+
+def worker_init_fn(worker_id: Any) -> None:
+    '''
+    Worker initialization function for DataLoader to ensure each worker has a unique seed.
+
+    Args:
+        worker_id (Any): The worker ID.
+    '''
+    worker_seed = torch.initial_seed() % 2**32
+    np.random.seed(worker_seed)
+    random.seed(worker_seed)
+
+def create_generator() -> torch.Generator:
+    '''
+    Creates a random number generator with the global seed.
+
+    Returns:
+        torch.Generator: The configured random number generator.
+    '''
+    import codon
+    seed = codon.__seed__ if hasattr(codon, '__seed__') and codon.__seed__ is not None else 42
+    return torch.Generator().manual_seed(seed)

codon/utils/theta.py

@@ -0,0 +1,55 @@
+import math
+from dataclasses import dataclass
+
+@dataclass
+class ValidateRoPEConfig:
+    '''
+    Configuration result for RoPE/Sinusoidal validation.
+
+    Attributes:
+        is_passed (bool): Whether the configuration is valid.
+        info (str): Assessment status.
+        suggested_base (float): The recommended base value.
+    '''
+    is_passed: bool
+    info: str
+    suggested_base: float
+
+def validate_rope_config(max_len: int, base: float) -> ValidateRoPEConfig:
+    '''
+    Validates if the RoPE/Sinusoidal base is sufficient for the given maximum length.
+
+    Args:
+        max_len (int): The maximum sequence length.
+        base (float): The base value for RoPE.
+
+    Returns:
+        ValidateRoPEConfig: The validation result containing status and recommendation.
+    '''
+    max_period = 2 * math.pi * base
+    
+    if max_len <= 8192:
+        recommended = 10000.0
+    else:
+        # base = 10000 * (scaling_factor ^ 1.1)
+        scaling_factor = max_len / 4096
+        recommended = 10000.0 * (scaling_factor ** 1.1)
+        
+    if recommended > 1000000:
+        recommended = round(recommended / 1000000) * 1000000.0
+    elif recommended > 100000:
+        recommended = round(recommended / 100000) * 100000.0
+    else:
+        recommended = round(recommended / 10000) * 10000.0
+
+    if max_period < max_len:
+        return ValidateRoPEConfig(is_passed=False, info='critical_low', suggested_base=recommended)
+    
+    elif max_period < max_len * 2:
+        return ValidateRoPEConfig(is_passed=False, info='low', suggested_base=recommended)
+    
+    elif max_period > max_len * 100:
+        return ValidateRoPEConfig(is_passed=True, info='too_high', suggested_base=recommended)
+    
+    else:
+        return ValidateRoPEConfig(is_passed=True, info='optimal', suggested_base=recommended)

codon/utils/token.py

@@ -0,0 +1,276 @@
+import json
+import os
+import zipfile
+from dataclasses import dataclass
+from tokenizers  import Tokenizer, pre_tokenizers, decoders
+from tokenizers  import normalizers
+from tokenizers.models   import BPE
+from tokenizers.trainers import BpeTrainer
+
+from transformers import PreTrainedTokenizerFast
+
+from typing import Union, Optional
+
+@dataclass
+class TokenizerTrainerResult:
+    '''
+    Result of the tokenizer trainer creation.
+
+    Attributes:
+        tokenizer (Tokenizer): The configured tokenizer instance.
+        trainer (BpeTrainer): The configured BPE trainer instance.
+    '''
+    tokenizer: Tokenizer
+    trainer: BpeTrainer
+
+
+core_tokens = ['[unk]', '[pad]', '[sep]']
+chat_tokens = [
+    '[im_start]', '[im_end]',
+    '[system]', '[user]', '[model]', '[tool]', '[train]',
+    '[interruption]', '[fim]',
+]
+reasoning_tokens = ['[cot_start]', '[cot_end]', '[verification]', '[solution]']
+code_tokens = ['[fim_pre]', '[fim_mid]', '[fim_suf]']
+tool_tokens = ['[tool_start]', '[tool_name]', '[tool_args]', '[tool_end]']
+
+multimodal_tokens = [
+    '[image_start]', '[image_end]', '[audio_start]', '[audio_end]', 
+    '[video_start]', '[video_end]'
+]
+
+base_special_tokens = (
+    core_tokens + 
+    chat_tokens + 
+    reasoning_tokens + 
+    code_tokens + 
+    tool_tokens + 
+    multimodal_tokens
+)
+
+base_special_tokens += [f'[unused_{i}]' for i in range(len(base_special_tokens), 64)]
+base_special_tokens += [f'[mask_{i}]' for i in range(32)]
+
+chat_template = (
+    "{% for message in messages %}"
+        "{{ '[im_start]' }}"
+
+        "{% if message['role'] == 'fim' %}"
+            "{{ '[fim]' }}"
+            "{{ '[fim_pre]' + message['prefix'] + '[fim_suf]' + message['suffix'] + '[fim_mid]' }}"
+            
+            "{% if message['middle'] %}"
+                "{{ message['middle'] + '[im_end]' }}"
+            "{% endif %}"
+            
+        "{% else %}"
+            
+            "{% if message['role'] in ['system', 'instruction'] %}"
+                "{{ '[system]' }}"
+            "{% elif message['role'] == 'user' %}"
+                "{{ '[user]' }}"
+            "{% elif message['role'] in ['assistant', 'model'] %}"
+                "{{ '[model]' }}"
+            "{% elif message['role'] == 'tool' %}"
+                "{{ '[tool]' }}"
+            "{% elif message['role'] == 'train' %}"
+                "{{ '[train]' }}"
+            "{% else %}"
+                "{{ message['role'] }}"
+            "{% endif %}"
+
+            "{{ '\n' }}"
+
+            "{% set thought_content = message['thought'] or message['reasoning_content'] %}"
+            "{% if thought_content %}"
+                "{{ '[cot_start]' + thought_content + '[cot_end]\n' }}"
+            "{% else %}"
+                "{{ '[cot_start][cot_end]\n' }}"
+            "{% endif %}"
+            
+            "{% if message['content'] is defined and message['content'] is not none %}"
+                "{% if message['content'] is string %}"
+                    "{{ message['content'] }}"
+                "{% else %}"
+                    "{% for item in message['content'] %}"
+                        "{% if item['type'] == 'text' %}"
+                            "{{ item['text'] }}"
+                        "{% elif item['type'] == 'image' %}"
+                            "{{ '[image_start][image_end]' }}"
+                        "{% elif item['type'] == 'audio' %}"
+                            "{{ '[audio_start][audio_end]' }}"
+                        "{% elif item['type'] == 'video' %}"
+                            "{{ '[video_start][video_end]' }}"
+                        "{% endif %}"
+                    "{% endfor %}"
+                "{% endif %}"
+            "{% endif %}"
+            
+            "{% if message['tool_calls'] is defined and message['tool_calls'] %}"
+                "{% for tool_call in message['tool_calls'] %}"
+                    "{{ '[tool_start][tool_name]' + tool_call.function.name + '[tool_args]' + tool_call.function.arguments + '[tool_end]' }}"
+                "{% endfor %}"
+            "{% endif %}"
+            
+            "{{ '[im_end]\n' }}"
+        "{% endif %}"
+    "{% endfor %}"
+    
+    "{% if add_generation_prompt %}"
+        "{% if messages[-1]['role'] != 'fim' %}"
+            "{{ '[im_start][model]\n' }}"
+            "{% if enable_thinking is defined and enable_thinking %}"
+                "{{ '[cot_start]' }}"
+            "{% elif enable_thinking is defined and not enable_thinking %}"
+                "{{ '[cot_start][cot_end]\n' }}"
+            "{% endif %}"
+        "{% endif %}"
+    "{% endif %}"
+)
+
+
+def create_tokenizer_trainer(
+    unk_token: str='[unk]',
+    vocab_size: int=32000,
+    special_tokens: list[str]=base_special_tokens
+) -> TokenizerTrainerResult:
+    '''
+    Creates a BPE Tokenizer trainer.
+
+    Configures and returns a tokenizer trainer object for training BPE (Byte-Pair Encoding) models.
+    The trainer is pre-configured with NFKC normalization, digit splitting, and byte-level pre-tokenization.
+
+    Args:
+        unk_token (str): Identifier for unknown tokens. Defaults to '[unk]'.
+        vocab_size (int): Target vocabulary size. Defaults to 32000.
+        special_tokens (list[str]): List of special tokens.
+            Defaults to base_special_tokens, including core, chat, reasoning, code, tool, and multimodal tokens.
+
+    Returns:
+        TokenizerTrainerResult: A dataclass containing the tokenizer and trainer instances.
+    '''
+    tokenizer = Tokenizer(BPE(unk_token=unk_token))
+
+    tokenizer.normalizer = normalizers.NFKC()
+
+    tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
+        pre_tokenizers.Digits(individual_digits=True),
+        pre_tokenizers.ByteLevel(
+            add_prefix_space=False,
+            use_regex=True
+        )
+    ])
+
+    tokenizer.decoder = decoders.ByteLevel()
+
+    trainer = BpeTrainer(
+        vocab_size=vocab_size,
+        special_tokens=special_tokens,
+        initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
+        max_token_length=32,
+        min_frequency=10
+    )
+
+    return TokenizerTrainerResult(tokenizer=tokenizer, trainer=trainer)
+
+
+class PackedTokenizer:
+    def __init__(self, tokenizer: Optional[Union[Tokenizer, str]]):
+        self._tokenizer: Optional[Tokenizer] = None
+        self._fast_tokenizer: Optional[PreTrainedTokenizerFast] = None
+        self.config = {}
+        self.template = chat_template
+
+        if isinstance(tokenizer, str):
+            self.load(tokenizer)
+        elif isinstance(tokenizer, Tokenizer):
+            self._tokenizer = tokenizer
+            self.config = {
+                'unk_token': '[unk]',
+                'pad_token': '[pad]',
+                'bos_token': '[im_start]',
+                'eos_token': '[im_end]',
+            }
+            self._update_fast_tokenizer()
+
+    def _update_fast_tokenizer(self) -> None:
+        '''
+        Updates the cached PreTrainedTokenizerFast instance.
+        '''
+        if self._tokenizer is None:
+            self._fast_tokenizer = None
+            return
+
+        self._fast_tokenizer = PreTrainedTokenizerFast(
+            tokenizer_object=self._tokenizer,
+            unk_token=self.config.get('unk_token', '[unk]'),
+            pad_token=self.config.get('pad_token', '[pad]'),
+            bos_token=self.config.get('bos_token', '[im_start]'),
+            eos_token=self.config.get('eos_token', '[im_end]'),
+            chat_template=self.template,
+            clean_up_tokenization_spaces=False
+        )
+
+    @property
+    def tokenizer(self) -> Tokenizer:
+        if self._tokenizer is None:
+            raise ValueError("Tokenizer is not loaded.")
+        return self._tokenizer
+    
+    @property
+    def fast_tokenizer(self) -> PreTrainedTokenizerFast:
+        if self._fast_tokenizer is None:
+            raise ValueError('Tokenizer is not loaded.')
+        return self._fast_tokenizer
+    
+    def save(self, path: str) -> 'PackedTokenizer':
+        if self._tokenizer is None:
+            raise ValueError("No tokenizer to save.")
+
+        with zipfile.ZipFile(path, 'w') as z:
+            # Save tokenizer.json
+            z.writestr('tokenizer.json', self._tokenizer.to_str())
+            
+            # Save tokenizer_config.json
+            z.writestr('tokenizer_config.json', json.dumps(self.config, indent=2))
+            
+            # Save chat_template.jinja
+            z.writestr('chat_template.jinja', self.template)
+            
+        return self
+    
+    def load(self, path: str) -> 'PackedTokenizer':
+        if not os.path.exists(path):
+            raise FileNotFoundError(f"File not found: {path}")
+
+        with zipfile.ZipFile(path, 'r') as z:
+            file_list = z.namelist()
+            
+            # Helper to find file in zip (ignoring directory prefix if any)
+            def find_file(name):
+                for f in file_list:
+                    if f == name or f.endswith(f'/{name}'):
+                        return f
+                return None
+
+            # Load tokenizer.json
+            tokenizer_file = find_file('tokenizer.json')
+            if tokenizer_file:
+                tokenizer_json = z.read(tokenizer_file).decode('utf-8')
+                self._tokenizer = Tokenizer.from_str(tokenizer_json)
+            else:
+                raise ValueError("tokenizer.json not found in zip file")
+
+            # Load tokenizer_config.json
+            config_file = find_file('tokenizer_config.json')
+            if config_file:
+                config_json = z.read(config_file).decode('utf-8')
+                self.config = json.loads(config_json)
+
+            # Load chat_template.jinja
+            template_file = find_file('chat_template.jinja')
+            if template_file:
+                self.template = z.read(template_file).decode('utf-8')
+
+        self._update_fast_tokenizer()
+        return self