accessible_worlds.aw_create_tokenizer

  1from transformers import AutoTokenizer, PreTrainedTokenizer
  2import orjson
  3import re
  4from pathlib import Path
  5from typing import override
  6
  7from .aw_vocabulary import AwVocabulary
  8
  9from .aw_utils import (
 10    split_and_preserve_names,
 11    load_json,
 12    safe_save
 13)
 14
 15import numpy as np
 16
 17def _normalize_word_spacing( text: str ) -> str :
 18
 19    # add space before opening quotes
 20    text = re.sub(r'"(?=[A-Za-z0-9])', '" ', text)
 21
 22    return text
 23
 24def _save_map_back( 
 25    remap_table : list[int] | np.ndarray, 
 26    map_back : dict[tuple,int], 
 27    filepath : str | Path 
 28) -> None :
 29
 30    data = {
 31        "aw_to_gemma" : [ int( id ) for id in remap_table ],
 32        "gemma_to_aw" : {",".join(map(str, k)): v for k, v in map_back.items()}
 33    }
 34    safe_save( data=data, filepath=filepath )
 35
 36def _load_map_back( filepath : str | Path ) -> tuple[ np.ndarray, dict[tuple,int] ] :
 37    
 38    data = load_json( filepath )
 39    return ( 
 40        data[ "aw_to_gemma" ], 
 41        { tuple(int(x) for x in k.split(",")): v for k, v in data["gemma_to_aw"].items() }
 42    )
 43
 44class AwTokenizer( ) :
 45
 46    # Note: Since Gemma 4 uses BPE with separate " x" vs "x" tokens, 
 47    # we opt for Moses-style processing and detokenization, ensuring spaces 
 48    # exist so we always get the " x" version.
 49
 50    def __init__( 
 51        self,
 52        tokenizer_path : str | Path,
 53        data_dir : str | Path,
 54        lowercase : bool,
 55        output_dir : str | Path ) :
 56
 57        base_tokenizer = AutoTokenizer.from_pretrained( tokenizer_path )
 58
 59        vocabulary_path = Path( data_dir ) / f"vocabulary.json"
 60
 61        if not vocabulary_path.exists() :
 62            raise ValueError( f"Vocabulary {vocabulary_path} not found" )
 63
 64        world_path = Path( data_dir ) / "world.json"
 65
 66        if not world_path.exists() :
 67            raise ValueError( f"World configuration {world_path} not found" )
 68
 69        exposition_path = Path( data_dir ) / "exposition_settings.json"
 70
 71        if not exposition_path.exists() :
 72            raise ValueError( f"Exposition configuration {exposition_path} not found" )
 73
 74        data_file_paths = [  
 75            Path( data_dir ) / f"data_{i}.json" for i in range( 1000 )
 76        ]
 77
 78        data_file_paths = [ p for p in data_file_paths if p.exists() ]
 79
 80        if not data_file_paths :
 81            raise ValueError( f"No files found in data_dir {data_dir}" )
 82
 83        vocabulary = AwVocabulary.from_file( vocabulary_path )
 84        world      = load_json( world_path )
 85        exposition = load_json( exposition_path )
 86
 87        token_set = set()
 88
 89        current_file_idx = 0
 90        while current_file_idx < len( data_file_paths ) :
 91            data = load_json( data_file_paths[ current_file_idx ] )    
 92            for instance in data["instances"] :
 93                text = instance["completion"].lower() if lowercase else instance["completion"]
 94                text = " " + _normalize_word_spacing( text )
 95                tokens = base_tokenizer.encode( text )
 96                token_set.update( tokens )
 97            current_file_idx += 1
 98
 99        for token in token_set :
100            print( base_tokenizer.decode( [ token ] ) )
101
102        self.remap_table = np.array( list(token_set), dtype=int )
103        for idx, token_id in enumerate( token_set ) :
104            self.remap_table[ idx ] = token_id
105            
106        self.vocab = {
107            base_tokenizer.encode( self.remap_table[ base_id ] ) : id
108            for id, base_id in enumerate( self.remap_table )
109        }
110
111        self.map_back = {}
112        for idx in range( len( self.remap_table ) ) :
113            
114            self.map_back[ (self.remap_table[ idx ],) ] = idx
115            word = base_tokenizer.decode( [ self.remap_table[ idx ] ] )
116            
117            # Map all unsupported variants back to the space normalized lowercase variant
118            if len( word ) > 1 and word[ 0 ] == ' ' :
119                
120                capitialized =  ' ' + word[ 1 ].upper() + word[2:]
121                ct = base_tokenizer.encode( capitialized )
122                self.map_back[ tuple( ct ) ] = idx
123
124                capitialized_no_space = capitialized[1:]
125                ctns = base_tokenizer.encode( capitialized_no_space )
126                self.map_back[ tuple( ctns ) ] = idx
127
128                no_space = word[1:]
129                ns = base_tokenizer.encode( no_space )
130                self.map_back[ tuple( ns ) ] = idx
131        
132        # _save_map_back( 
133        #     remap_table=self.remap_table, 
134        #     map_back=self.map_back, 
135        #     filepath=Path(output_dir)/"map_back.json" 
136        # )
137
138
139if __name__ == "__main__":
140
141    sw_tokenizer = AwTokenizer( 
142        tokenizer_path="google/gemma-4-26B-A4B-it",
143        data_dir="/media/ubuntu/VisData/a-machine/data/simpler_worlds/aw-v1",
144        lowercase=True,
145        output_dir="/media/ubuntu/VisData/a-machine/data/simpler_worlds/aw-v1/tokenizer"
146    )