@@ -300,77 +300,77 @@ def predict(self, docs: Iterable[Doc]) -> List[str]:
300300 for i , doc in enumerate (docs ):
301301 sentences = [s for s in doc .sents ]
302302 if len (doc ) > 0 :
303- # Looping through each sentence and each entity
304- # This may go wrong if there are entities across sentences - which shouldn't happen normally.
305- for sent_index , sent in enumerate (sentences ):
306- if sent .ents :
307- # get n_neightbour sentences, clipped to the length of the document
308- start_sentence = max (0 , sent_index - self .n_sents )
309- end_sentence = min (
310- len (sentences ) - 1 , sent_index + self .n_sents
311- )
312- start_token = sentences [start_sentence ].start
313- end_token = sentences [end_sentence ].end
314- sent_doc = doc [start_token :end_token ].as_doc ()
315- # currently, the context is the same for each entity in a sentence (should be refined)
316- xp = self .model .ops .xp
317- if self .incl_context :
318- sentence_encoding = self .model .predict ([sent_doc ])[0 ]
319- sentence_encoding_t = sentence_encoding .T
320- sentence_norm = xp .linalg .norm (sentence_encoding_t )
321- for ent in sent .ents :
322- entity_count += 1
323- if ent .label_ in self .labels_discard :
324- # ignoring this entity - setting to NIL
325- final_kb_ids .append (self .NIL )
326- else :
327- candidates = self .get_candidates (self .kb , ent )
328- if not candidates :
329- # no prediction possible for this entity - setting to NIL
330- final_kb_ids .append (self .NIL )
331- elif len (candidates ) == 1 :
332- # shortcut for efficiency reasons: take the 1 candidate
333- # TODO: thresholding
334- final_kb_ids .append (candidates [0 ].entity_ )
335- else :
336- random .shuffle (candidates )
337- # set all prior probabilities to 0 if incl_prior=False
338- prior_probs = xp .asarray (
339- [c .prior_prob for c in candidates ]
340- )
341- if not self .incl_prior :
342- prior_probs = xp .asarray (
343- [0.0 for _ in candidates ]
344- )
345- scores = prior_probs
346- # add in similarity from the context
347- if self .incl_context :
348- entity_encodings = xp .asarray (
349- [c .entity_vector for c in candidates ]
350- )
351- entity_norm = xp .linalg .norm (
352- entity_encodings , axis = 1
353- )
354- if len (entity_encodings ) != len (prior_probs ):
355- raise RuntimeError (
356- Errors .E147 .format (
357- method = "predict" ,
358- msg = "vectors not of equal length" ,
359- )
360- )
361- # cosine similarity
362- sims = xp .dot (
363- entity_encodings , sentence_encoding_t
364- ) / (sentence_norm * entity_norm )
365- if sims .shape != prior_probs .shape :
366- raise ValueError (Errors .E161 )
367- scores = (
368- prior_probs + sims - (prior_probs * sims )
303+ # Looping through each entity (TODO: rewrite)
304+ for ent in doc .ents :
305+ sent = ent .sent
306+ sent_index = sentences .index (sent )
307+ assert sent_index >= 0
308+ # get n_neightbour sentences, clipped to the length of the document
309+ start_sentence = max (0 , sent_index - self .n_sents )
310+ end_sentence = min (
311+ len (sentences ) - 1 , sent_index + self .n_sents
312+ )
313+ start_token = sentences [start_sentence ].start
314+ end_token = sentences [end_sentence ].end
315+ sent_doc = doc [start_token :end_token ].as_doc ()
316+ # currently, the context is the same for each entity in a sentence (should be refined)
317+ xp = self .model .ops .xp
318+ if self .incl_context :
319+ sentence_encoding = self .model .predict ([sent_doc ])[0 ]
320+ sentence_encoding_t = sentence_encoding .T
321+ sentence_norm = xp .linalg .norm (sentence_encoding_t )
322+ entity_count += 1
323+ if ent .label_ in self .labels_discard :
324+ # ignoring this entity - setting to NIL
325+ final_kb_ids .append (self .NIL )
326+ else :
327+ candidates = self .get_candidates (self .kb , ent )
328+ if not candidates :
329+ # no prediction possible for this entity - setting to NIL
330+ final_kb_ids .append (self .NIL )
331+ elif len (candidates ) == 1 :
332+ # shortcut for efficiency reasons: take the 1 candidate
333+ # TODO: thresholding
334+ final_kb_ids .append (candidates [0 ].entity_ )
335+ else :
336+ random .shuffle (candidates )
337+ # set all prior probabilities to 0 if incl_prior=False
338+ prior_probs = xp .asarray (
339+ [c .prior_prob for c in candidates ]
340+ )
341+ if not self .incl_prior :
342+ prior_probs = xp .asarray (
343+ [0.0 for _ in candidates ]
344+ )
345+ scores = prior_probs
346+ # add in similarity from the context
347+ if self .incl_context :
348+ entity_encodings = xp .asarray (
349+ [c .entity_vector for c in candidates ]
350+ )
351+ entity_norm = xp .linalg .norm (
352+ entity_encodings , axis = 1
353+ )
354+ if len (entity_encodings ) != len (prior_probs ):
355+ raise RuntimeError (
356+ Errors .E147 .format (
357+ method = "predict" ,
358+ msg = "vectors not of equal length" ,
369359 )
370- # TODO: thresholding
371- best_index = scores .argmax ().item ()
372- best_candidate = candidates [best_index ]
373- final_kb_ids .append (best_candidate .entity_ )
360+ )
361+ # cosine similarity
362+ sims = xp .dot (
363+ entity_encodings , sentence_encoding_t
364+ ) / (sentence_norm * entity_norm )
365+ if sims .shape != prior_probs .shape :
366+ raise ValueError (Errors .E161 )
367+ scores = (
368+ prior_probs + sims - (prior_probs * sims )
369+ )
370+ # TODO: thresholding
371+ best_index = scores .argmax ().item ()
372+ best_candidate = candidates [best_index ]
373+ final_kb_ids .append (best_candidate .entity_ )
374374 if not (len (final_kb_ids ) == entity_count ):
375375 err = Errors .E147 .format (
376376 method = "predict" , msg = "result variables not of equal length"
0 commit comments