1from transformers import AutoTokenizer
2import argparse
3from pathlib import Path
4from unidecode import unidecode
5
6from .aw_utils import (
7 load_json,
8 safe_save,
9 in_dictionary
10)
11
12if __name__ == "__main__":
13
14 this_dir = Path(__file__).parent
15 data = load_json( this_dir / "data/names.json" )
16
17 parser = argparse.ArgumentParser()
18
19 parser.add_argument( "--tokenizer_path", type=str )
20 parser.add_argument( "--output_dir", type=str )
21
22 args = parser.parse_args()
23
24 tokenizer_path = args.tokenizer_path
25 output_dir = Path( args.output_dir )
26
27 if not output_dir.exists() :
28 raise ValueError( f"Output directory {output_dir} doesn't exist." )
29
30 tokenizer = AutoTokenizer.from_pretrained( tokenizer_path )
31
32 s_token = tokenizer.encode( 's' )
33 quote_token = tokenizer.encode( '\"' )
34
35 results = {
36 "female" : [],
37 "male" : []
38 }
39
40 for country in data :
41 for gender in data[ country ] :
42 for name in data[ country ][ gender ] :
43
44 ascii_name = unidecode( name )
45
46 orig_letters = sum(1 for c in name if c.isalpha())
47 ascii_letters = sum(1 for c in ascii_name if c.isalpha())
48
49 if orig_letters == 0 or ascii_letters != orig_letters :
50 continue
51
52 name = ascii_name.strip()
53
54 if ( name[0].islower()
55 or in_dictionary( name.lower() )
56 or in_dictionary( name.lower() + "s" )
57 or " " in name
58 or "-" in name
59 or name[-1] == "s" ) :
60 continue
61
62 after_quote = tokenizer.encode( '\"' + name, add_special_tokens=False )
63 after_quote_plural = tokenizer.encode( '\"' + name + 's', add_special_tokens=False )
64
65 after_quote_and_space = tokenizer.encode( '\" ' + name, add_special_tokens=False )
66 after_quote_and_space_plural = tokenizer.encode( '\" ' + name + 's', add_special_tokens=False )
67
68 tokens = tokenizer.encode( name, add_special_tokens=False )
69 tokens_plural = tokenizer.encode( name + 's', add_special_tokens=False )
70
71 tokens_space = tokenizer.encode( ' ' + name, add_special_tokens=False )
72 tokens_plural_space = tokenizer.encode( ' ' + name + 's', add_special_tokens=False )
73
74 valid_without_space = False
75 valid_with_space = False
76 valid_after_quote_and_space = False
77
78 if (
79 len( tokens_space ) == 1
80 and len( tokens_plural_space ) == 2
81 and tokens_plural_space[1] == s_token[ 0 ]
82 and tokens_space[ 0 ] == tokens_plural_space[ 0 ]
83 and len( after_quote_and_space ) == 2
84 and len( after_quote_and_space_plural ) == 3
85 and after_quote_and_space[0] == quote_token[ 0 ]
86 and after_quote_and_space_plural[0] == quote_token[ 0 ]
87 and after_quote_and_space_plural[2] == s_token[ 0 ]
88 and after_quote_and_space[ 1 ] == after_quote_and_space_plural[ 1 ]
89 ) :
90 valid_with_space = True
91
92 if (
93 len( tokens ) == 1
94 and len( tokens_plural ) == 2
95 and tokens_plural[1] == s_token[ 0 ]
96 and tokens[ 0 ] == tokens_plural[ 0 ]
97 and len( after_quote ) == 2
98 and len( after_quote_plural ) == 3
99 and after_quote[0] == quote_token[ 0 ]
100 and after_quote_plural[0] == quote_token[ 0 ]
101 and after_quote_plural[2] == s_token[ 0 ]
102 and after_quote[ 1 ] == after_quote_plural[ 1 ]
103 ) :
104 valid_without_space = True
105
106 if valid_with_space and valid_without_space and name not in results[ gender ] :
107 results[ gender ].append( name )
108 print( name )
109
110 safe_save( data=results, filepath=output_dir/"single_token_names.json" )