accessible_worlds.aw_vocabulary

  1from collections import Counter
  2from pathlib import Path
  3
  4import numpy as np
  5
  6from .aw_utils import (
  7    safe_save, 
  8    load_json, 
  9    synonyms_wsd,
 10    get_word_forms
 11)
 12
 13class AwVocabulary :
 14
 15    def __init__( 
 16        self,
 17        core_things_and_concepts : dict[str,list[str]],
 18        moral_concepts : dict[str,list[str]],
 19        word_counter : Counter | None = None,
 20        core_integration_counter  : Counter | None = None,
 21        moral_integration_counter : Counter | None = None,
 22        all_core_word_forms : set[str] | None = None,
 23        vocab_word_forms : dict[str, list[str] ] | None = None,
 24        all_vocab_word_forms : set[str] | None = None,
 25        vocab_names : set[str] | None = None ) :
 26
 27        self.core_things_and_concepts = core_things_and_concepts
 28        self.moral_concepts = moral_concepts
 29
 30        if word_counter is None :
 31            self.word_counter = Counter()
 32        else :
 33            self.word_counter = word_counter
 34
 35        if core_integration_counter is None :
 36            self.core_integration_counter = Counter( {
 37                integration : 0
 38                for _, category in self.core_things_and_concepts.items() for integration in category
 39            } )
 40        else :
 41            self.core_integration_counter = core_integration_counter
 42
 43        if moral_integration_counter is None :
 44            self.moral_integration_counter = Counter( {
 45                integration : 0
 46                for _, category in self.moral_concepts.items() for integration in category
 47            } )
 48        else :
 49            self.moral_integration_counter = moral_integration_counter
 50
 51        if all_core_word_forms is None :
 52            all_words = set()
 53            for element in list( self.core_integration_counter.keys() ) + list( self.moral_integration_counter.keys() ) :
 54                all_words.update( element.split() )
 55
 56            word_forms = get_word_forms( all_words )
 57            self.all_core_word_forms = set( [ w for _, forms in word_forms.items() for w in forms ] )
 58        else :
 59            self.all_core_word_forms = all_core_word_forms
 60
 61        if all_vocab_word_forms is None :
 62            self.all_vocab_word_forms = set()
 63        else :
 64            self.all_vocab_word_forms = all_vocab_word_forms
 65
 66        if vocab_word_forms is None :
 67            self.vocab_word_forms = dict()
 68        else :
 69            self.vocab_word_forms = vocab_word_forms
 70
 71        if vocab_names is None :
 72            self.vocab_names = set() 
 73        else :
 74            self.vocab_names = vocab_names
 75
 76    @classmethod
 77    def from_file( 
 78        cls, 
 79        filepath : str | Path ) :
 80
 81        data = load_json( filepath )
 82        return cls(
 83            core_things_and_concepts=data['core_things_and_concepts'],
 84            moral_concepts=data[ 'moral_concepts' ],
 85            word_counter=Counter( data['word_counter'] ),
 86            core_integration_counter=Counter( data['core_integration_counter'] ),
 87            moral_integration_counter=Counter( data['moral_integration_counter'] ), 
 88            all_core_word_forms=set( data['all_core_word_forms'] ),
 89            vocab_word_forms=data['vocab_word_forms'],
 90            all_vocab_word_forms=set( data['all_vocab_word_forms'] ),
 91            vocab_names=set(data['vocab_names'])
 92        )
 93
 94    def to_file( self, filepath : str | Path ) : 
 95        
 96        data =  {
 97            "core_things_and_concepts" : self.core_things_and_concepts,
 98            "moral_concepts" : self.moral_concepts,
 99            "word_counter" : self.word_counter,
100            "core_integration_counter" : self.core_integration_counter,
101            "moral_integration_counter" : self.moral_integration_counter,
102            "all_core_word_forms"  : list( self.all_core_word_forms ),
103            "all_vocab_word_forms" : list( self.all_vocab_word_forms ),
104            "vocab_word_forms" : self.vocab_word_forms,
105            "vocab_names" : list( self.vocab_names )
106        }
107
108        safe_save( 
109            data=data, 
110            filepath=filepath )
111
112    def update_words( self, words ) :
113
114        self.word_counter.update( words )
115        new_words = list( set( [ w for w in words if not w in self.vocab_word_forms ] ) )
116        new_word_forms = get_word_forms( new_words )
117        self.vocab_word_forms.update( { w : list(forms) for w, forms in new_word_forms.items() } )
118        all_new_forms = [ wf for w, forms in new_word_forms.items() for wf in forms ]
119        self.all_vocab_word_forms.update( all_new_forms )
120
121        self.vocab_names.update( [ word for word in words if word and word[0].isupper() ] )
122
123    def update_core_integrations( self, integrations : list[str] ) :
124        self.core_integration_counter.update( integrations )
125
126    def update_moral_integrations( self, integrations : list[str] ) :
127        self.moral_integration_counter.update( integrations )
128
129    def sample_least_gounded( self, min_k : int, max_k : int, np_rng ) -> list[str] :
130
131        # no boosting single characters, contraction tokens, or names
132        filtered_counts = {k: v for k, v in self.word_counter.items() if ( len(k) > 1 and "'" not in k and k[0].islower() ) }
133
134        if not filtered_counts :
135            return []
136        
137        counts = np.array( list( filtered_counts.values() ), dtype=float )
138        counts = np.where( counts==0, 1.0, counts )
139        probs = ( 1.0 / counts**2 )
140        probs = probs / np.sum( probs )
141
142        k = min( np_rng.integers( min_k, max_k+1 ), len( self.word_counter ) )
143
144        return np_rng.choice(  
145            list( filtered_counts.keys() ), 
146            size=k, 
147            replace=False, 
148            p=probs ).tolist()
149
150    def _sample_integrations( self, counter : Counter, np_rng, min_k : int, max_k : int ) -> list[str] :
151        
152        k = min( len( counter ), np_rng.integers( min_k, max_k+1 ) )
153        chosen = []
154        min_count = min( counter.values() )
155        min_items = [ str(key) for key, count in counter.items()  if count == min_count ]
156        
157        while len( chosen ) < k :
158
159            if len( min_items ) >= k :
160                chosen += np_rng.choice( min_items, size=k, replace=False ).tolist()
161                break
162            else :
163                chosen += min_items 
164                k -= len( min_items )
165                min_count = min( [ val for val in counter.values() if val > min_count ] )
166                min_items = [ key for key, count in counter.items()  if count == min_count ]
167
168        return chosen
169
170    def sample_core_integrations( 
171        self, 
172        np_rng, 
173        min_k : int, 
174        max_k : int,
175        min_extra_creature_integrations : int,
176        max_extra_creature_integrations : int ) -> list[str] :
177
178        core_integrations = self._sample_integrations( 
179            self.core_integration_counter, 
180            np_rng, 
181            min_k=min_k, 
182            max_k=max_k 
183        )
184
185        n_extra_creature_integrations = np_rng.integers( 
186            min_extra_creature_integrations, 
187            max_extra_creature_integrations+1 
188        )
189
190        if max_extra_creature_integrations > 0 :
191            
192            extra_creature_integrations = np_rng.choice(  
193                self.core_things_and_concepts[ 'representative creatures' ],
194                size=n_extra_creature_integrations,
195                replace=False
196            ).tolist()
197
198            core_integrations = list( set( core_integrations + extra_creature_integrations ) )
199
200        return core_integrations 
201
202    def sample_moral_integrations( self, np_rng, min_k : int, max_k : int ) -> list[str] :
203        return self._sample_integrations( self.moral_integration_counter, np_rng, min_k=min_k, max_k=max_k )
204
205    def synonyms( self, word, context : str ) -> list[str] :
206        all_synonyms = synonyms_wsd( word, context )
207        return [ s for s in all_synonyms if s in self.word_counter ]
208
209    def get_word_counts( self, words : list[str] ) -> dict[str, int] : 
210        return {
211            word: self.word_counter.get( word, 0 )
212            for word in set( words )
213        }
214
215    def get_words( self ) -> list[str] : 
216        return list( self.word_counter.keys() ) 
217
218    def get_names( self ) -> set[str] : 
219        return self.vocab_names
220
221    def size( self ) -> int :
222        return len( self.word_counter )
223
224    def has_word( self, word : str ) -> bool : 
225        return word in self.word_counter
226
227    def is_core_word( self, word : str ) -> bool : 
228        return ( 
229            word in self.core_integration_counter or 
230            word in self.moral_integration_counter
231        )
232
233    def is_form_of_core_word( self, word : str ) -> bool : 
234        return word in self.all_core_word_forms
235
236    def is_form_of_vocab_word( self, word : str ) -> bool : 
237        return word in self.all_vocab_word_forms