@@ -58,6 +58,7 @@ def __init__(self, cui2info: dict[str, CUIInfo],
5858 def get_context_tokens (self , entity : MutableEntity , doc : MutableDocument ,
5959 size : int ,
6060 per_doc_valid_token_cache : 'PerDocumentTokenCache' ,
61+ fill_centre_tokens : bool = True ,
6162 ) -> tuple [list [MutableToken ],
6263 list [MutableToken ],
6364 list [MutableToken ]]:
@@ -83,17 +84,21 @@ def get_context_tokens(self, entity: MutableEntity, doc: MutableDocument,
8384 per_doc_valid_token_cache [tkn ]]
8485 # Reverse because the first token should be the one closest to center
8586 tokens_left .reverse ()
86- tokens_center : list [MutableToken ] = list (
87- cast (Iterable [MutableToken ], entity ))
87+ if fill_centre_tokens :
88+ tokens_center : list [MutableToken ] = list (
89+ cast (Iterable [MutableToken ], entity ))
90+ else :
91+ tokens_center = []
8892 _right_tokens = doc [end_ind + 1 :end_ind + 1 + size ]
8993 tokens_right = [tkn for tkn in _right_tokens if
9094 per_doc_valid_token_cache [tkn ]]
9195
9296 return tokens_left , tokens_center , tokens_right
9397
94- def _tokens2vecs (self , tokens : Sequence [Union [MutableToken , str ]]
95- ) -> Iterable [np .ndarray ]:
96- for step , tkn in enumerate (tokens ):
98+ def _tokens2vecs (self , tokens : Sequence [Union [MutableToken , str ]],
99+ step_start : int = 0
100+ ) -> Iterable [np .ndarray ]:
101+ for step , tkn in enumerate (tokens , start = step_start ):
97102 lower = tkn .lower () if isinstance (tkn , str ) else tkn .base .lower
98103 if lower not in self .vocab :
99104 continue
@@ -137,26 +142,55 @@ def get_context_vectors(self, entity: MutableEntity,
137142 """
138143 vectors : dict [str , np .ndarray ] = {}
139144
140- context_vector_sizes = self .config .context_vector_sizes
141- for context_type , window_size in context_vector_sizes .items ():
142- tokens_left , tokens_center , tokens_right = self .get_context_tokens (
143- entity , doc , window_size , per_doc_valid_token_cache )
144-
145- values : list [np .ndarray ] = []
146- # Add left
147- values .extend (self ._tokens2vecs (tokens_left ))
148-
149- if not self .config .context_ignore_center_tokens :
150- # Add center
151- values .extend (
152- self ._preprocess_center_tokens (cui , tokens_center ))
153-
154- # Add right
155- values .extend (self ._tokens2vecs (tokens_right ))
156-
145+ # Sort ascending so each iteration is a superset of the previous
146+ sorted_contexts = sorted (
147+ self .config .context_vector_sizes .items (), key = lambda x : x [1 ])
148+
149+ prev_left : list [MutableToken ] = []
150+ prev_right : list [MutableToken ] = []
151+ # Accumulated weighted vecs from previous (smaller) windows,
152+ # excluding center (center is the same for all window sizes)
153+ prev_left_vecs : list [np .ndarray ] = []
154+ prev_right_vecs : list [np .ndarray ] = []
155+
156+ # Center is identical for all window sizes, only compute once
157+ if not self .config .context_ignore_center_tokens :
158+ tokens_center = list (
159+ cast (Iterable [MutableToken ], entity ))
160+ center_vecs = list (
161+ self ._preprocess_center_tokens (cui , tokens_center ))
162+ else :
163+ center_vecs = []
164+
165+ for context_type , window_size in sorted_contexts :
166+ tokens_left , _ , tokens_right = self .get_context_tokens (
167+ entity , doc , window_size , per_doc_valid_token_cache ,
168+ fill_centre_tokens = False )
169+
170+ # New outer tokens only — the inner ones were already processed
171+ # NOTE: left hand tokens are in order of closest first, which is why
172+ # we're slicing from the start of the list
173+ new_left = tokens_left [len (prev_left ):]
174+ new_right = tokens_right [len (prev_right ):]
175+
176+ # step_start for new left tokens: they are further from centre
177+ # so their step index is
178+ # len(tokens_left) - len(new_left) ... len(tokens_left)-1
179+ # i.e. the new tokens are the outermost, highest-step ones
180+ new_left_vecs = list (self ._tokens2vecs (
181+ new_left , step_start = len (prev_left )))
182+ new_right_vecs = list (self ._tokens2vecs (
183+ new_right , step_start = len (prev_right )))
184+
185+ prev_left_vecs = new_left_vecs + prev_left_vecs
186+ prev_right_vecs = prev_right_vecs + new_right_vecs
187+ prev_left = tokens_left
188+ prev_right = tokens_right
189+
190+ values = prev_left_vecs + center_vecs + prev_right_vecs
157191 if values :
158- value = np .average (values , axis = 0 )
159- vectors [ context_type ] = value
192+ vectors [ context_type ] = np .average (values , axis = 0 )
193+
160194 return vectors
161195
162196 def similarity (self , cui : str , entity : MutableEntity , doc : MutableDocument ,
0 commit comments