From d41d875186da336abeb5b6f05e5cf4676c5855b8 Mon Sep 17 00:00:00 2001
From: svlandeg <svlandeg@github.com>
Date: Wed, 27 Mar 2024 16:37:45 +0100
Subject: [PATCH] restore entities of sample in initialization

---
 spacy/pipeline/entity_linker.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/spacy/pipeline/entity_linker.py b/spacy/pipeline/entity_linker.py
index 3f70475fd..0581d51ab 100644
--- a/spacy/pipeline/entity_linker.py
+++ b/spacy/pipeline/entity_linker.py
@@ -284,9 +284,11 @@ class EntityLinker(TrainablePipe):
         nO = self.kb.entity_vector_length
         doc_sample = []
         vector_sample = []
+        orig_ents = []
         for eg in islice(get_examples(), 10):
             doc = eg.x
             if self.use_gold_ents:
+                orig_ents.append(doc.ents)
                 ents, _ = eg.get_aligned_ents_and_ner()
                 doc.ents = ents
             doc_sample.append(doc)
@@ -313,6 +315,10 @@ class EntityLinker(TrainablePipe):
         if not has_annotations:
             # Clean up dummy annotation
             doc.ents = []
+        if self.use_gold_ents:
+            assert len(doc_sample) == len(orig_ents)
+            for doc, orig_ent in zip(doc_sample, orig_ents):
+                doc.ents = orig_ent
 
     def batch_has_learnable_example(self, examples):
         """Check if a batch contains a learnable example.