accessible_worlds.aw_name_stratification

  1import argparse
  2from pathlib import Path
  3import re
  4import random
  5import string
  6from collections import Counter
  7
  8from .aw_utils import (
  9    load_json,
 10    safe_save
 11)
 12
 13from .aw_vocabulary import AwVocabulary
 14
 15def random_string() -> str:
 16    return "".join(random.choices(string.ascii_letters, k=10))
 17
 18def find_names( passage: str, name_set: set[str] ) -> set[str]:
 19    found = set()
 20    for name in name_set:
 21        pattern = r'\b' + re.escape(name) + r's?\b'
 22        if re.search(pattern, passage):
 23            found.add(name)
 24    return found
 25
 26def replace_name( passage: str, previous_name: str, new_name: str ) -> tuple[str, int, int]:
 27
 28    pattern = r'\b' + re.escape(previous_name) + r'(s)?\b'
 29    
 30    singular_count = 0
 31    plural_count = 0
 32    
 33    def do_replacement( m ):
 34        nonlocal singular_count, plural_count
 35        if m.group(1):
 36            plural_count += 1
 37            return new_name + 's'
 38        else:
 39            singular_count += 1
 40            return new_name
 41            
 42    new_passage = re.sub( pattern, do_replacement, passage )
 43
 44    return new_passage, singular_count, plural_count
 45
 46if __name__ == "__main__":
 47
 48    # ---------------------------------------------------------------
 49
 50    parser = argparse.ArgumentParser()
 51    parser.add_argument( "--data_dir",   type=str )
 52    parser.add_argument( "--name_pool_path", type=str )
 53    parser.add_argument( "--random_seed", type=int, default=32 )
 54    args = parser.parse_args()
 55
 56    data_dir = Path( args.data_dir )
 57    name_pool_path  = Path( args.name_pool_path )
 58    random_seed = args.random_seed
 59
 60    # ---------------------------------------------------------------------------------------
 61
 62    # Validate files
 63
 64    random.seed( random_seed )
 65
 66    data_file_paths = [  ( data_dir / f"data_{i}.json", i ) for i in range( 1000 ) ]
 67    data_file_paths = [  ( p, idx ) for p, idx in data_file_paths if p.exists() ] 
 68
 69    vocab_path = data_dir / f"vocabulary.json"
 70
 71    if not data_file_paths :
 72        raise ValueError( f"No data files found in {data_dir}" )
 73
 74    if not vocab_path.exists() :
 75        raise ValueError( f"No vocabulary in {data_dir}" )
 76
 77    vocabulary = AwVocabulary.from_file( vocab_path )
 78
 79    name_pool = {}
 80    if name_pool_path.exists() :
 81        name_pool = load_json( name_pool_path )
 82    else :
 83        raise ValueError( "No name pool given" )
 84
 85    # ---------------------------------------------------------------------------------------
 86    # Validate name pool and initialize data structures to track names
 87
 88    existing_vocab_names = set( [ name for name in vocabulary.vocab_names if name[-1] != "s" ] )
 89    existing_vocab_names = existing_vocab_names | set( [ name[:-1] for name in vocabulary.vocab_names if name[-1] == "s" ] )
 90
 91    previous_names = set( name_pool[ "previous_names" ][ "male" ] + name_pool[ "previous_names" ][ "female" ] )
 92    
 93    previous_name_genders = {
 94        name : gender 
 95        for gender in name_pool[ "previous_names" ] for name in name_pool[ "previous_names" ][gender]
 96    }
 97
 98    if previous_names != existing_vocab_names :
 99        raise ValueError( "Previous names in pool must be the same as vocab names in vocabulary, up to plurality." )
100
101    new_names = set( name_pool[ "new_names" ][ "male" ] + name_pool[ "new_names" ][ "female" ] ) 
102    
103    new_name_genders = {
104        name : gender 
105        for gender in name_pool[ "new_names" ] for name in name_pool[ "new_names" ][gender]
106    }
107
108    # Must have at least as many new per-gender as previous, to avoid running out of candidates as we replace
109    # in a given instance
110    previous_gender_counts = Counter( previous_name_genders.values() )
111    new_gender_counts      = Counter( new_name_genders.values() )
112
113    for gender, count in previous_gender_counts.items() :
114        if new_gender_counts.get( gender, 0 ) < count :
115            raise ValueError( f"Must have at least as many new names as precious per-gender"  )
116
117    # -------------------------------------------------------------------------------------------
118    # Data structures to count for stratification and new total occurences
119
120    counts_general = {
121        name : 0 
122        for name in new_names
123    }
124
125    occurence_singular = {
126        name : 0 
127        for name in new_names
128    }
129
130    occurence_plural = {
131        name : 0 
132        for name in new_names
133    }
134
135    counts_per_moral_intg = {}
136
137    # ---------------------------------------------------------------------------------------
138    # Apply stratified name replacements incrementally
139
140    for data_path in data_file_paths :
141
142        data = load_json( data_path[0] )
143        
144        # for each instance in the file
145        for instance in data[ "instances" ] :
146
147            # existing content potentially containinig names            
148            passage = instance[ 'completion' ]
149            moral_intgs = instance[ "moral_integrations" ]
150            recursive_integration = instance[ "recursive_integration" ]
151            prompt = instance[ "prompt" ]
152
153            # passage names occurences are the only that matter for stratification
154            passage_names = find_names( passage, previous_names )
155
156            # initialize name per moral integration count if hasn't ocurred yet
157            for intg in moral_intgs :
158                if intg not in counts_per_moral_intg :
159                    counts_per_moral_intg[ intg ] = {
160                        name : 0 
161                        for name in new_names
162                    }
163
164            # for tracking what names have already been used and which have been replaced to avoid collision
165            used_in_instance = set()
166            replaced_in_passage = set()
167
168            # for storing temporary name replacements in order to avoid collisions when previous and new names aren't disjoint
169            replacements : list[tuple[str,str]] = []
170
171            # for all names found in the passage
172            for name in passage_names : 
173                
174                # get the gender of the name
175                name_gender = previous_name_genders[ name ]
176
177                # Use general count if no moral integrations
178                if not moral_intgs :
179                    candidates = [
180                        ( new_name, count )
181                        for new_name, count in counts_general.items() if ( 
182                            new_name_genders[ new_name ] == name_gender and 
183                            not new_name in used_in_instance
184                        )
185                    ]
186
187                # map candidates based on moral integration counts if moral integrations used on instance
188                else : 
189                    candidates = [
190                        ( new_name, count )
191                        for intg in moral_intgs for new_name, count in counts_per_moral_intg[intg].items() if ( 
192                            new_name_genders[ new_name ] == name_gender and 
193                            not new_name in used_in_instance
194                        )
195                    ]
196
197                def score( new_name ):
198                    counts = [ count for _, count in candidates ]
199                    return ( max( counts ), sum( counts ) )  # minimize worst case, then total
200
201                # select randomly among the best candidates
202                min_score = min( score( name ) for name, _ in candidates )
203                best_candidates = [name for name, _ in candidates if score( name ) == min_score]
204                replacement_name = random.choice(best_candidates)
205
206                # track names used for replacement in this instance
207                used_in_instance.add( replacement_name )
208                replaced_in_passage.add( name )
209                
210                # replace with a temporary unique id to avoid collision with non-disjoint new and previous name sets 
211                temporary_rep = f"{random_string()}{len(replacements)}"
212                passage, singular_count, plural_count = replace_name( passage=passage, previous_name=name, new_name=temporary_rep )
213                recursive_integration, _, _ = replace_name( passage=recursive_integration, previous_name=name, new_name=temporary_rep )
214                prompt, _, _ = replace_name( passage=prompt, previous_name=name, new_name=temporary_rep )
215
216                # keep track of tempory unique id assignment to actual replacement names
217                replacements.append( ( temporary_rep, replacement_name ) )
218
219                # keep track of total occurences for the actual replacement names
220                occurence_singular[ replacement_name ] += singular_count
221                occurence_plural[   replacement_name ]   += plural_count
222
223                # keeping track of how many passages the names appear in
224                counts_general[ replacement_name ] += 1
225
226                # keeping track of appearences in passage by moral integration
227                for intg in moral_intgs :
228                    counts_per_moral_intg[ intg ][ replacement_name ] += 1
229
230            # Handle names that were in a prompt, or recursive integration, that weren't in the completion
231            for name in find_names( f"{prompt} {recursive_integration}", previous_names - replaced_in_passage ) :
232
233                # Get name candidates
234                # Statificatin doesn't matter for names that didn't end up in the completion
235                # but they must be unused so far in this instance
236                candidates = [
237                    new_name for new_name in new_names 
238                    if new_name_genders[ new_name ] ==  previous_name_genders[ name ] and new_name not in used_in_instance
239                ]
240                replacement_name = random.choice( candidates )
241                
242                # track names used
243                used_in_instance.add( replacement_name )
244
245                # replace with a temporary unique id to avoid collision with non-disjoint new and previous name sets 
246                temporary_rep = f"{random_string()}{len(replacements)}"
247                recursive_integration, _, _ = replace_name( passage=recursive_integration, previous_name=name, new_name=temporary_rep )
248                prompt, _, _ = replace_name( passage=prompt, previous_name=name, new_name=temporary_rep )
249
250                # keep track of tempory unique id assignment to actual replacement names
251                replacements.append( ( temporary_rep, replacement_name ) )
252
253            # Finally replace the temporay unique ids with the corresponding new names
254            for replacement_name in replacements :
255
256                passage, _, _ = replace_name( 
257                    passage=passage, 
258                    previous_name=replacement_name[0], 
259                    new_name=replacement_name[1] 
260                )
261                
262                recursive_integration, _, _ = replace_name( 
263                    passage=recursive_integration, 
264                    previous_name=replacement_name[0], 
265                    new_name=replacement_name[1] 
266                
267                )
268                prompt, _, _ = replace_name( 
269                    passage=prompt, 
270                    previous_name=replacement_name[0], 
271                    new_name=replacement_name[1] 
272                )
273
274            # update the instance with the new content
275            instance[ 'completion' ] = passage
276            instance[ 'recursive_integration' ] = recursive_integration
277            instance[ 'prompt' ] = prompt
278
279        # Save the updated data file with new name 
280        safe_save( data=data, filepath=data_dir / f"data_{data_path[1]}_stratified.json" )
281
282    # Update the vocabulary name list
283    vocabulary.vocab_names = { name for name in new_names } | { f"{name}s" for name in new_names } 
284    
285    # Remove previous names from vocabulary word tracking data structures
286    for name in previous_name_genders : 
287        vocabulary.word_counter.pop( name, None )
288        vocabulary.word_counter.pop( f"{name}s", None )
289        vocabulary.all_core_word_forms.discard( name )
290        vocabulary.all_core_word_forms.discard( f"{name}s" )
291        vocabulary.all_vocab_word_forms.discard( name )
292        vocabulary.all_vocab_word_forms.discard( f"{name}s" )
293        vocabulary.vocab_word_forms.pop( name, None )
294        vocabulary.vocab_word_forms.pop( f"{name}s", None )
295
296    # Add the new names the vocabulary data structures
297    for name in new_name_genders : 
298        vocabulary.all_core_word_forms.add( name )
299        vocabulary.all_core_word_forms.add( f"{name}s" )
300        vocabulary.vocab_word_forms[ name      ] = [ name, f"{name}s" ]
301        vocabulary.vocab_word_forms[ f"{name}s"] = [ name, f"{name}s" ]
302
303    # Update the new total counts for the names
304    for name, count in occurence_singular.items() :
305        vocabulary.word_counter[ name ] = count
306
307    for name, count in occurence_plural.items() :
308        vocabulary.word_counter[ f"{name}s" ] = count
309
310    # Save the updated vocabulary
311    vocabulary.to_file( filepath=data_dir / f"vocabulary_stratified.json" )