accessible_worlds.aw_utils

  1import orjson
  2import os
  3from typing import Any
  4from pathlib import Path
  5import re
  6from collections import defaultdict
  7
  8from nltk.wsd import lesk
  9from nltk.corpus import wordnet
 10import lemminflect
 11from word_forms import word_forms as wf
 12
 13import en_core_web_sm
 14_nlp = en_core_web_sm.load()
 15
 16import enchant 
 17
 18_broker = enchant.Broker()
 19_broker.set_ordering("en_GB", "nuspell,hunspell,aspell")
 20_broker.set_ordering("en_US", "nuspell,hunspell,aspell")
 21
 22_gb_dictionary = _broker.request_dict( "en_GB" )
 23_us_dictionary = _broker.request_dict( "en_US" )
 24
 25_added_dictionary = set()
 26_this_dir = Path(__file__).parent
 27
 28_blacklisted_words = set()
 29_blacklisted_names = set()
 30
 31with open( _this_dir / "data/added_words.txt", 'r') as file:
 32    _lines = file.readlines()
 33    for line in _lines :
 34        w = line.strip()
 35        if w :
 36            _added_dictionary.add( w )
 37
 38with open( _this_dir / "data/blacklisted_words.txt", 'r') as file:
 39    _lines = file.readlines()
 40    for line in _lines :
 41        w = line.strip()
 42        if w :
 43            _blacklisted_words.add( w )
 44
 45with open( _this_dir / "data/blacklisted_names.txt", 'r') as file:
 46    _lines = file.readlines()
 47    for line in _lines :
 48        name = line.strip()
 49        if name  :
 50            _blacklisted_names.add( name )
 51
 52def in_dictionary( word : str ) -> bool :
 53    return (
 54        ( _gb_dictionary.check( word ) or _us_dictionary.check( word ) or word in _added_dictionary ) and
 55        ( not word in _blacklisted_words ) 
 56    )
 57 
 58def _rgb(r, g, b):
 59    return f"\033[38;2;{r};{g};{b}m"
 60
 61PINK   = _rgb(255, 105, 180)
 62PURPLE = _rgb(200, 150, 255)
 63BLUE   = _rgb(135, 206, 250)
 64GREEN  = _rgb(144, 238, 144)
 65YELLOW = _rgb(255, 255, 102)
 66ORANGE = _rgb(255, 178, 102)
 67TEAL   = _rgb(135, 255, 240)
 68RESET  = "\033[0m"
 69
 70def safe_save( data : dict[str,Any], filepath : str | Path ) -> None :
 71	temp_path = f"{filepath}.tmp"
 72	try:
 73		with open(temp_path, 'wb') as f:
 74			f.write(orjson.dumps(data, option=orjson.OPT_INDENT_2 | orjson.OPT_SERIALIZE_NUMPY ))
 75			f.flush()
 76			os.fsync(f.fileno())
 77		os.replace(temp_path, filepath)
 78	
 79	except Exception as e:
 80		print(f"An error occurred while saving the file: {e}")
 81		if os.path.exists(temp_path):
 82			try:
 83				os.remove(temp_path)
 84			except OSError as cleanup_error:
 85				print(f"Failed to clean up temporary file '{temp_path}': {cleanup_error}")
 86
 87def load_json( filepath : str | Path ) -> dict[str, Any] :
 88	with open(filepath, 'rb') as f:
 89		return orjson.loads(f.read())
 90
 91def split_sentences(passage: str):
 92    return [ s.strip() for s in re.split(r'(?<=[.?!])', passage ) if s.strip() ]
 93
 94def split_words( passage : str ) -> list[str] :
 95    return re.findall( r"[a-z]+(?:-[a-z]+)*|'s", passage.lower() )
 96
 97def split_paragraphs( passage : str ) -> list[str] :
 98    return re.split(r'\n\s*\n', passage.strip())
 99
100def spacy_to_wordnet(tag: str) -> str | None:
101    
102    if not tag:
103        return None
104
105    if tag.startswith('V'): return wordnet.VERB
106    if tag.startswith('N'): return wordnet.NOUN
107    if tag.startswith('J'): return wordnet.ADJ
108    if tag.startswith('R'): return wordnet.ADV
109
110    return None
111
112def synonyms_wsd(word: str, context: str) -> list[str]:
113
114    doc = _nlp( context )
115    target_token = None
116
117    for token in doc:
118        if token.text.lower() == word.lower():
119            target_token = token
120            break
121
122    if target_token is None:
123        return []
124
125    tag = target_token.tag_
126
127    wn_pos = spacy_to_wordnet(tag)
128
129    if wn_pos is None:
130        return []
131
132    base_word = target_token.lemma_.lower()
133
134    context_tokens = [t.text.lower() for t in doc]
135    best_synset = lesk(context_tokens, base_word, pos=wn_pos)
136
137    if best_synset is None:
138        return []
139
140    results = set()
141
142    for lemma in best_synset.lemmas():
143        
144        lemma_name = lemma.name().lower().replace("_", " ")
145
146        if lemma_name == base_word or " " in lemma_name:
147            continue
148
149        all_inflections = lemminflect.getAllInflections(lemma_name)
150        if tag in all_inflections:
151            results.add(all_inflections[tag][0])
152        else:
153            results.add(lemma_name)  # fallback to base form
154
155    return list(results)
156
157def serialize_np_rng_state( state: dict ) -> dict:
158    s = state.copy()
159    s['state'] = {
160        'state': str(state['state']['state']),
161        'inc':   str(state['state']['inc']),
162    }
163    return s
164
165def deserialize_np_rng_state( state: dict ) -> dict:
166    s = state.copy()
167    s['state'] = {
168        'state': int(state['state']['state']),
169        'inc':   int(state['state']['inc']),
170    }
171    return s
172
173def get_word_forms( words: list[str] ) -> dict[str, set[str]]:
174    result = {}
175    for word in words:
176        forms = wf.get_word_forms(word)
177        merged_forms = set().union(*forms.values())
178        merged_forms.add(word)
179        result[word] = merged_forms
180    return result
181
182def preprocess_passage( passage: str ) -> str :
183
184    replace_map = {
185        "—"  : ",",
186        "”"  : '"', 
187        "“"  : '"',
188        "‟"  : '"',
189        "″"  : '"', 
190        "’"  : "'", 
191        "‘"  : "'", 
192        "‛"  : "'", 
193        "′"  : "'", 
194        "＇" : "'", 
195        "＂" : '"', 
196        "ʼ"  : "'",
197        "\t" : " ", 
198        "\xa0" : " ",
199        "—": ",",
200        "―": ",",
201        "‐": "-",
202        "‑": "-",
203        "‒": "-",
204        "−": "-",
205        "－": "-",
206        "–" : "-"
207    }
208
209    passage = passage.translate( str.maketrans( replace_map ) )
210    passage = re.sub( r"-{2,}", "-", passage )
211    passage = re.sub( r"\s*,",  ",", passage )
212    passage = re.sub( r" {2,}", " ", passage )
213
214    return passage
215
216def split_and_preserve_names(
217    passage: str,
218    vocabulary,
219    known_names : set[str],
220    known_names_only : bool = True ) -> tuple[list[str], bool]:
221    
222    results = []
223    is_start_of_sentence = True
224    last_end = 0
225    is_valid = True
226
227    # No double dash, sometimes used as em-dash
228    if "--" in passage :
229        is_valid = False
230        print( f"Double dash." )
231
232    quotes_open = False
233
234    invalid_chars = re.search(r'[^A-Za-z0-9\s.,!?;:\'"\-]', passage)
235
236    if invalid_chars:
237        is_valid = False
238        print( f"Bad characters found {invalid_chars}" )
239
240    potential_name_counts = defaultdict(int)
241
242    for match in re.finditer(r"[A-Za-z]+(?:-[A-Za-z]+)*|'(?:s|t|ll|d|re|ve|m)", passage):
243        
244        token = match.group()
245
246        gap = passage[last_end:match.start()]
247
248        # Matches punctuation
249        if re.search(r'[.!?]', gap):
250            is_start_of_sentence = True
251
252        # detect quotes
253        if re.search(r'"', gap) :
254            if quotes_open :
255                quotes_open = False 
256            else:
257                quotes_open = True
258                is_start_of_sentence = True
259
260        if token in [ "'s", "'t", "'ll", "'d", "'re", "'ve", "'m" ]:
261            results.append(token)
262            last_end = match.end()
263            continue
264
265        is_capitalized = token[0].isupper()
266
267        is_word = vocabulary.is_form_of_core_word( token.lower() ) or in_dictionary( token.lower() )
268
269        # word in dictionary that is capitalized and not start of sentence and not 'I'
270        if token != 'I' and is_capitalized and not is_start_of_sentence and is_word :
271            is_valid = False
272            print( f"Bad token {token}" )
273
274        elif token.lower() in _blacklisted_words or token.lower() in _blacklisted_names :
275            is_valid = False
276            print( f"Blacklisted word or name {token}" )
277
278        # no hyphens
279        elif '-' in token :
280
281            joined = token.replace( '-', '' ).lower()
282            tsplit = token.split( '-' )
283
284            if ( 
285                ( vocabulary.is_form_of_core_word( joined ) or in_dictionary( joined ) ) 
286            and ( not joined in _blacklisted_words )
287            ) :
288                results.append( joined )
289
290            elif ( len( tsplit ) == 2 and
291                ( vocabulary.is_form_of_core_word( tsplit[0].lower() ) or in_dictionary( tsplit[0].lower() ) ) and
292                ( vocabulary.is_form_of_core_word( tsplit[1].lower() ) or in_dictionary( tsplit[1].lower() ) ) and 
293            not ( tsplit[0].lower() in _blacklisted_words ) and 
294            not ( tsplit[1].lower() in _blacklisted_words )
295            ):  
296                results.append( tsplit[0].lower() )
297                results.append( tsplit[1].lower() )
298            else :
299                is_valid = False
300                print( f"Bad token {token}" )
301                
302        elif ( ( len(token) == 1 and is_start_of_sentence and token not in [ "A", "I" ] )
303            or ( len(token) == 1 and not is_start_of_sentence and token not in [ "I", "a" ] )
304        ) :
305            is_valid = False
306            print( f"Bad token {token}" )
307
308        # Non-name that isn't in the dictionary
309        elif not is_capitalized and not is_start_of_sentence and not is_word :
310            is_valid = False
311            print( f"Bad token {token}" )
312
313        else : 
314            if is_capitalized and not is_word :
315                potential_name_counts[ token ] += 1
316                results.append( token )
317            else:
318                results.append(token.lower())
319
320        is_start_of_sentence = False
321        last_end = match.end()
322
323    # merge plural with singular names
324    merged_potential_name_counts = defaultdict( lambda: [0, False, False] )
325
326    for potential_name, count in potential_name_counts.items() :
327        
328        if potential_name.endswith( 's' ) :
329            base = potential_name[:-1]
330            if base in known_names or base in potential_name_counts :
331                merged_potential_name_counts[ base ][0] += count
332                merged_potential_name_counts[ base ][2] = True
333            # If not a known name with without the s, or form without the s in the passage
334            # then we assume will treat it as plural (and it will fail later since no non-plural variant found)
335            else :
336                merged_potential_name_counts[ potential_name ][0] += count
337                merged_potential_name_counts[ potential_name ][2] = True
338        else :
339            merged_potential_name_counts[ potential_name ][0] += count
340            merged_potential_name_counts[ potential_name ][1] = True
341
342    # Names usually appear more than once in a passage. 
343    # This guards mainly against inconsistently spelled names or misspelled non-names at the start of a sentence
344    for potential_name, ( count, has_singular, has_plural ) in merged_potential_name_counts.items() :
345
346        if known_names_only and potential_name not in known_names :
347            is_valid = False 
348            print( f"Non-known name {potential_name}" )
349
350        elif potential_name not in known_names:
351
352            if has_plural and not has_singular :
353                is_valid = False 
354                print( f"Non-known names with s at end must have non-plural variant without {potential_name}" )
355
356            if re.search(r'(e|es)$', potential_name) :
357                is_valid = False
358                print(f"Ambiguous e/es-ending name {potential_name}")
359
360            if count < 2 :
361                is_valid = False 
362                print( f"Single count non-known name {potential_name}" )
363                
364            if potential_name in _blacklisted_names :
365                is_valid = False 
366                print( f"Blacklisted name {potential_name}" )
367
368    if not is_valid :
369        print( f"Invalid passage: {passage}" )
370
371    return results, is_valid
372
373if __name__ == "__main__":
374
375    # data = load_json( "/media/ubuntu/VisData/a-machine/data/simpler_worlds/sw-g26b-1.44m/sw-g26b-1.44m_0.json")
376    # for e in data[ 'instances' ] :
377    #     words, is_valid = split_and_preserve_names( e['completion' ] )
378    #     for word in words : 
379    #         ...
380
381    # def _split_words( passage : str ) -> list[str] :
382    #     return re.findall( r"[a-z]+(?:[a-z]+)*|'s", passage.lower() )
383
384    # # w = _split_words( "This is a test.\nHere is another word.\n...And I don't know\n\nBut lets see." )
385    # data_path = Path( "/media/ubuntu/VisData/a-machine/data/stories/_tokenizer_train_cache/simplestories.jsonl" )
386
387    # from collections import Counter
388    # counter = Counter()
389
390    # with data_path.open("r", encoding="utf-8") as f:
391    #     for line in f:
392    #         line = line.replace( r'\n', ' ' )
393    #         words = _split_words( line )
394    #         counter.update( words )
395
396    # print( [ key for key, c in  counter.items() if c < 2 ] )
397
398    # total_words = sum( [ c for key, c in counter.items() ] )
399
400    # print( f"Total words {total_words}\n" )
401
402    # total_words = sum( [ 1 for key, c in counter.items() ] )
403    # print( f"Total unique words {total_words}\n" )
404
405    # print( "Frequency of words that individually occure less than i times" )
406    # for i in range( 2, 100 ) :
407    #     print( f"count(w) < {i} : {sum( [ c for key, c in counter.items() if c < i ] )}" )
408
409    # print( "\nNumber of unique words that individually occure less than i times" )
410    # for i in range( 2, 100 ) :
411    #     print( f"count(w) < {i} : {sum( [ 1 for key, c in counter.items() if c < i ] )}" )
412
413    # print( "\nFrequency of words that are not in the dictionary" )
414    # print( f"{sum( [ c for key, c in counter.items() if not in_dictionary( key.lower() ) ] )}\n" )
415
416    # print( "Number of unique words that are not in the dictionary" )
417    # print( f"{sum( [ 1 for key, c in counter.items() if not in_dictionary( key.lower() ) ] )}\n" )
418
419    # n_total = sum( counter.values() )
420    # print( len( counter ) )
421    # for i in range( 1000 ) :
422    #     print( f"count(w) < {i} : {100*sum( [ c for key, c in counter.items() if c < i ] ) / n_total}%" )
423
424    # print( len(counter) )
425    # l1=[ w for w, c in counter.items() if c == 1 ]
426    # print( f"{1: } {len(l1)}" )
427    # print( l1 )
428
429    print( in_dictionary( "elara" ) )
430    print( in_dictionary( "tulay" ) )
431    print( in_dictionary( "olivia" ) )
432    print( in_dictionary( "sofia" ) )
433    print( in_dictionary( "alice" ) )
434
435    print( in_dictionary( "leo" ) )
436    print( in_dictionary( "kael" ) )
437    print( in_dictionary( "minh" ) )
438    print( in_dictionary( "yuto" ) )
439    print( in_dictionary( "aarav" ) )
440    print( in_dictionary( "liam" ) )
441
442    print()
443
444    print( in_dictionary( "elaras" ) )
445    print( in_dictionary( "tulays" ) )
446    print( in_dictionary( "olivias" ) )
447    print( in_dictionary( "sofias" ) )
448    print( in_dictionary( "alices" ) )
449    print( in_dictionary( "kaels" ) )
450    print( in_dictionary( "minhs" ) )
451    print( in_dictionary( "yutos" ) )
452    print( in_dictionary( "aaravs" ) )
453    print( in_dictionary( "liams" ) )
454    print( in_dictionary( "leos" ) )
455
456    print( in_dictionary( "elar" ) )
457    print( in_dictionary( "elars" ) )
458    print( in_dictionary( "kumal" ) )
459    print( in_dictionary( "kumals" ) )
460    print()
461
462    print( in_dictionary( "gray" ) )
463    print( in_dictionary( "color" ) )
464    print( in_dictionary( "grey" ) )
465    print( in_dictionary( "colour" ) )
466    print( in_dictionary( "interconnectedness" ) )
467    print()
468
469    print( in_dictionary( "sunday" ) )
470    print( in_dictionary( "Sunday" ) )
471
472    print( in_dictionary( "sunday" ) )
473    print( in_dictionary( "Sunday" ) )
474
475    # print( _blacklisted_words )
476    # print( _blacklisted_names )
477
478    # common_names = load_json( _this_dir / "data/common_names.json" )
479    # safe_save( data=common_names, filepath=_this_dir / "data/common_names.json" )