We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[Bug]: Avg Pooling in the Entity Linker
### Describe the bug

A runtime error is raised upon prediction when using "average" as the pooling operation in the Entity Linker

### To Reproduce

```python
from flair.data import Corpus
from flair.datasets import NEL_ENGLISH_TWEEKI
from flair.embeddings import TransformerWordEmbeddings
from flair.models import EntityLinker

corpus: Corpus = NEL_ENGLISH_TWEEKI(sample_missing_splits=False)

embeddings = TransformerWordEmbeddings(
    model="distilbert-base-uncased",
    fine_tune=True,
)

entity_linker = EntityLinker(
    embeddings=embeddings,
    label_dictionary=corpus.make_label_dictionary(label_type="nel"),
    label_type="nel",
    pooling_operation="average",
)

entity_linker.predict(corpus.train[0])
```


### Expected behaivor

The Entity Linker should be able to perform average pooling without any issues, like in the case of other pooling options.

### Logs and Stack traces

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In [1], line 20
      8 embeddings = TransformerWordEmbeddings(
      9     model="distilbert-base-uncased",
     10     fine_tune=True,
     11 )
     13 entity_linker = EntityLinker(
     14     embeddings=embeddings,
     15     label_dictionary=corpus.make_label_dictionary(label_type="nel"),
     16     label_type="nel",
     17     pooling_operation="average",
     18 )
---> 20 entity_linker.predict(corpus.train[0])

File ~/projects/flair_forked/flair/nn/model.py:826, in DefaultClassifier.predict(self, sentences, mini_batch_size, return_probabilities_for_all_classes, verbose, label_name, return_loss, embedding_storage_mode)
    824 # pass data points through network and decode
    825 data_point_tensor = self._encode_data_points(batch, data_points)
--> 826 scores = self.decoder(data_point_tensor)
    827 scores = self._mask_scores(scores, data_points)
    829 # if anything could possibly be predicted

File ~/miniforge3/envs/flair/lib/python3.9/site-packages/torch/nn/modules/module.py:1190, in Module._call_impl(self, *input, **kwargs)
   1186 # If we don't have any hooks, we want to skip the rest of the logic in
   1187 # this function, and just call forward.
   1188 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1189         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1190     return forward_call(*input, **kwargs)
   1191 # Do not call functions when jit is used
   1192 full_backward_hooks, non_full_backward_hooks = [], []

File ~/miniforge3/envs/flair/lib/python3.9/site-packages/torch/nn/modules/linear.py:114, in Linear.forward(self, input)
    113 def forward(self, input: Tensor) -> Tensor:
--> 114     return F.linear(input, self.weight, self.bias)

RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x4 and 768x650)
```

### Screenshots

_No response_

### Additional Context

_No response_

### Environment

#### Versions:
##### Flair
0.11.3
##### Pytorch
1.13.0
##### Transformers
4.24.0
#### GPU
False

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `flair/models/entity_linker_model.py`
Content:
```
1 import logging
2 import re
3 from functools import lru_cache
4 from pathlib import Path
5 from typing import Any, Callable, Dict, List, Optional, Set, Union
6 from unicodedata import category
7 
8 import torch
9 
10 import flair.embeddings
11 import flair.nn
12 from flair.data import Dictionary, Sentence, Span
13 from flair.file_utils import cached_path
14 
15 log = logging.getLogger("flair")
16 
17 
18 class CandidateGenerator:
19     """
20     Given a string, the CandidateGenerator returns possible target classes as candidates.
21     """
22 
23     def __init__(self, candidates: Union[str, Dict], backoff: bool = True):
24         # internal candidate lists of generator
25         self.mention_to_candidates_map: Dict = {}
26 
27         # load Zelda candidates if so passed
28         if isinstance(candidates, str) and candidates.lower() == "zelda":
29             zelda_path: str = "https://flair.informatik.hu-berlin.de/resources/datasets/zelda"
30             zelda_candidates = cached_path(f"{zelda_path}/zelda_mention_entities_counter.pickle", cache_dir="datasets")
31             import pickle
32 
33             with open(zelda_candidates, "rb") as handle:
34                 mention_entities_counter = pickle.load(handle)
35 
36             # create candidate lists
37             candidate_lists = {}
38             for mention in mention_entities_counter:
39                 candidate_lists[mention] = list(mention_entities_counter[mention].keys())
40 
41             self.mention_to_candidates_map = candidate_lists
42 
43         elif isinstance(candidates, Dict):
44             self.mention_to_candidates_map = candidates
45 
46         # if lower casing is enabled, create candidate lists of lower cased versions
47         self.backoff = backoff
48         if self.backoff:
49             # create a new dictionary for lower cased mentions
50             lowercased_mention_to_candidates_map: Dict = {}
51 
52             # go through each mention and its candidates
53             for mention, candidates in self.mention_to_candidates_map.items():
54                 backoff_mention = self._make_backoff_string(mention)
55                 # check if backoff mention already seen. If so, add candidates. Else, create new entry.
56                 if backoff_mention in lowercased_mention_to_candidates_map:
57                     current_candidates = lowercased_mention_to_candidates_map[backoff_mention]
58                     lowercased_mention_to_candidates_map[backoff_mention] = set(current_candidates).union(candidates)
59                 else:
60                     lowercased_mention_to_candidates_map[backoff_mention] = candidates
61 
62             # set lowercased version as map
63             self.mention_to_candidates_map = lowercased_mention_to_candidates_map
64 
65     @lru_cache(maxsize=50000)
66     def _make_backoff_string(self, mention: str) -> str:
67         backoff_mention = mention.lower()
68         backoff_mention = "".join(ch for ch in backoff_mention if category(ch)[0] not in "P")
69         backoff_mention = re.sub(" +", " ", backoff_mention)
70         return backoff_mention
71 
72     def get_candidates(self, mention: str) -> Set[str]:
73         """Given a mention, this method returns a set of candidate classes"""
74         if self.backoff:
75             mention = self._make_backoff_string(mention)
76 
77         return set(self.mention_to_candidates_map[mention]) if mention in self.mention_to_candidates_map else set()
78 
79 
80 class EntityLinker(flair.nn.DefaultClassifier[Sentence, Span]):
81     """
82     Entity Linking Model
83     The model expects text/sentences with annotated entity mentions and predicts entities to these mentions.
84     To this end a word embedding is used to embed the sentences and the embedding of the entity mention goes through a linear layer to get the actual class label.
85     The model is able to predict '<unk>' for entity mentions that the model can not confidently match to any of the known labels.
86     """
87 
88     def __init__(
89         self,
90         embeddings: flair.embeddings.TokenEmbeddings,
91         label_dictionary: Dictionary,
92         pooling_operation: str = "first_last",
93         label_type: str = "nel",
94         candidates: Optional[CandidateGenerator] = None,
95         **classifierargs,
96     ):
97         """
98         Initializes an EntityLinker
99         :param embeddings: embeddings used to embed the words/sentences
100         :param label_dictionary: dictionary that gives ids to all classes. Should contain <unk>
101         :param pooling_operation: either 'average', 'first', 'last' or 'first&last'. Specifies the way of how text representations of entity mentions (with more than one word) are handled.
102         E.g. 'average' means that as text representation we take the average of the embeddings of the words in the mention. 'first&last' concatenates
103         the embedding of the first and the embedding of the last word.
104         :param label_type: name of the label you use.
105         """
106 
107         super(EntityLinker, self).__init__(
108             embeddings=embeddings,
109             label_dictionary=label_dictionary,
110             final_embedding_size=embeddings.embedding_length * 2
111             if pooling_operation == "first_last"
112             else embeddings.embedding_length,
113             **classifierargs,
114         )
115 
116         self.pooling_operation = pooling_operation
117         self._label_type = label_type
118 
119         cases: Dict[str, Callable[[Span, List[str]], torch.Tensor]] = {
120             "average": self.emb_mean,
121             "first": self.emb_first,
122             "last": self.emb_last,
123             "first_last": self.emb_firstAndLast,
124         }
125 
126         if pooling_operation not in cases:
127             raise KeyError('pooling_operation has to be one of "average", "first", "last" or "first_last"')
128 
129         self.aggregated_embedding = cases[pooling_operation]
130 
131         self.candidates = candidates
132 
133         self.to(flair.device)
134 
135     def emb_first(self, span: Span, embedding_names):
136         return span.tokens[0].get_embedding(embedding_names)
137 
138     def emb_last(self, span: Span, embedding_names):
139         return span.tokens[-1].get_embedding(embedding_names)
140 
141     def emb_firstAndLast(self, span: Span, embedding_names):
142         return torch.cat(
143             (span.tokens[0].get_embedding(embedding_names), span.tokens[-1].get_embedding(embedding_names)), 0
144         )
145 
146     def emb_mean(self, span, embedding_names):
147         return torch.mean(torch.cat([token.get_embedding(embedding_names) for token in span], 0), 0)
148 
149     def _get_data_points_from_sentence(self, sentence: Sentence) -> List[Span]:
150         return sentence.get_spans(self.label_type)
151 
152     def _filter_data_point(self, data_point: Sentence) -> bool:
153         return bool(data_point.get_labels(self.label_type))
154 
155     def _get_embedding_for_data_point(self, prediction_data_point: Span) -> torch.Tensor:
156         return self.aggregated_embedding(prediction_data_point, self.embeddings.get_names())
157 
158     def _get_state_dict(self):
159         model_state = {
160             **super()._get_state_dict(),
161             "word_embeddings": self.embeddings.save_embeddings(use_state_dict=False),
162             "label_type": self.label_type,
163             "label_dictionary": self.label_dictionary,
164             "pooling_operation": self.pooling_operation,
165             "loss_weights": self.weight_dict,
166             "candidates": self.candidates,
167         }
168         return model_state
169 
170     def _print_predictions(self, batch, gold_label_type):
171         lines = []
172         for datapoint in batch:
173             eval_line = f"\n{datapoint.to_original_text()}\n"
174 
175             for span in datapoint.get_spans(gold_label_type):
176                 symbol = "✓" if span.get_label(gold_label_type).value == span.get_label("predicted").value else "❌"
177                 eval_line += (
178                     f' - "{span.text}" / {span.get_label(gold_label_type).value}'
179                     f' --> {span.get_label("predicted").value} ({symbol})\n'
180                 )
181 
182             lines.append(eval_line)
183         return lines
184 
185     @classmethod
186     def _init_model_with_state_dict(cls, state, **kwargs):
187         # remap state dict for models serialized with Flair <= 0.11.3
188         import re
189 
190         state_dict = state["state_dict"]
191         for key in list(state_dict.keys()):
192             state_dict[re.sub("^word_embeddings\\.", "embeddings.", key)] = state_dict.pop(key)
193 
194         return super()._init_model_with_state_dict(
195             state,
196             embeddings=state.get("word_embeddings"),
197             label_dictionary=state.get("label_dictionary"),
198             label_type=state.get("label_type"),
199             pooling_operation=state.get("pooling_operation"),
200             loss_weights=state.get("loss_weights", {"<unk>": 0.3}),
201             candidates=state.get("candidates", None),
202             **kwargs,
203         )
204 
205     @property
206     def label_type(self):
207         return self._label_type
208 
209     def _mask_scores(self, scores: torch.Tensor, data_points: List[Span]):
210         if not self.candidates:
211             return scores
212 
213         masked_scores = -torch.inf * torch.ones(scores.size(), requires_grad=True, device=flair.device)
214 
215         for idx, span in enumerate(data_points):
216             # get the candidates
217             candidate_set = self.candidates.get_candidates(span.text)
218             # during training, add the gold value as candidate
219             if self.training:
220                 candidate_set.add(span.get_label(self.label_type).value)
221             candidate_set.add("<unk>")
222             indices_of_candidates = [self.label_dictionary.get_idx_for_item(candidate) for candidate in candidate_set]
223             masked_scores[idx, indices_of_candidates] = scores[idx, indices_of_candidates]
224 
225         return masked_scores
226 
227     @classmethod
228     def load(cls, model_path: Union[str, Path, Dict[str, Any]]) -> "EntityLinker":
229         from typing import cast
230 
231         return cast("EntityLinker", super().load(model_path=model_path))
232
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/flair/models/entity_linker_model.py b/flair/models/entity_linker_model.py
--- a/flair/models/entity_linker_model.py
+++ b/flair/models/entity_linker_model.py
@@ -144,7 +144,7 @@
         )
 
     def emb_mean(self, span, embedding_names):
-        return torch.mean(torch.cat([token.get_embedding(embedding_names) for token in span], 0), 0)
+        return torch.mean(torch.stack([token.get_embedding(embedding_names) for token in span], 0), 0)
 
     def _get_data_points_from_sentence(self, sentence: Sentence) -> List[Span]:
         return sentence.get_spans(self.label_type)

{"golden_diff": "diff --git a/flair/models/entity_linker_model.py b/flair/models/entity_linker_model.py\n--- a/flair/models/entity_linker_model.py\n+++ b/flair/models/entity_linker_model.py\n@@ -144,7 +144,7 @@\n         )\n \n     def emb_mean(self, span, embedding_names):\n-        return torch.mean(torch.cat([token.get_embedding(embedding_names) for token in span], 0), 0)\n+        return torch.mean(torch.stack([token.get_embedding(embedding_names) for token in span], 0), 0)\n \n     def _get_data_points_from_sentence(self, sentence: Sentence) -> List[Span]:\n         return sentence.get_spans(self.label_type)\n", "issue": "[Bug]: Avg Pooling in the Entity Linker\n### Describe the bug\r\n\r\nA runtime error is raised upon prediction when using \"average\" as the pooling operation in the Entity Linker\r\n\r\n### To Reproduce\r\n\r\n```python\r\nfrom flair.data import Corpus\r\nfrom flair.datasets import NEL_ENGLISH_TWEEKI\r\nfrom flair.embeddings import TransformerWordEmbeddings\r\nfrom flair.models import EntityLinker\r\n\r\ncorpus: Corpus = NEL_ENGLISH_TWEEKI(sample_missing_splits=False)\r\n\r\nembeddings = TransformerWordEmbeddings(\r\n    model=\"distilbert-base-uncased\",\r\n    fine_tune=True,\r\n)\r\n\r\nentity_linker = EntityLinker(\r\n    embeddings=embeddings,\r\n    label_dictionary=corpus.make_label_dictionary(label_type=\"nel\"),\r\n    label_type=\"nel\",\r\n    pooling_operation=\"average\",\r\n)\r\n\r\nentity_linker.predict(corpus.train[0])\r\n```\r\n\r\n\r\n### Expected behaivor\r\n\r\nThe Entity Linker should be able to perform average pooling without any issues, like in the case of other pooling options.\r\n\r\n### Logs and Stack traces\r\n\r\n```\r\n---------------------------------------------------------------------------\r\nRuntimeError                              Traceback (most recent call last)\r\nCell In [1], line 20\r\n      8 embeddings = TransformerWordEmbeddings(\r\n      9     model=\"distilbert-base-uncased\",\r\n     10     fine_tune=True,\r\n     11 )\r\n     13 entity_linker = EntityLinker(\r\n     14     embeddings=embeddings,\r\n     15     label_dictionary=corpus.make_label_dictionary(label_type=\"nel\"),\r\n     16     label_type=\"nel\",\r\n     17     pooling_operation=\"average\",\r\n     18 )\r\n---> 20 entity_linker.predict(corpus.train[0])\r\n\r\nFile ~/projects/flair_forked/flair/nn/model.py:826, in DefaultClassifier.predict(self, sentences, mini_batch_size, return_probabilities_for_all_classes, verbose, label_name, return_loss, embedding_storage_mode)\r\n    824 # pass data points through network and decode\r\n    825 data_point_tensor = self._encode_data_points(batch, data_points)\r\n--> 826 scores = self.decoder(data_point_tensor)\r\n    827 scores = self._mask_scores(scores, data_points)\r\n    829 # if anything could possibly be predicted\r\n\r\nFile ~/miniforge3/envs/flair/lib/python3.9/site-packages/torch/nn/modules/module.py:1190, in Module._call_impl(self, *input, **kwargs)\r\n   1186 # If we don't have any hooks, we want to skip the rest of the logic in\r\n   1187 # this function, and just call forward.\r\n   1188 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks\r\n   1189         or _global_forward_hooks or _global_forward_pre_hooks):\r\n-> 1190     return forward_call(*input, **kwargs)\r\n   1191 # Do not call functions when jit is used\r\n   1192 full_backward_hooks, non_full_backward_hooks = [], []\r\n\r\nFile ~/miniforge3/envs/flair/lib/python3.9/site-packages/torch/nn/modules/linear.py:114, in Linear.forward(self, input)\r\n    113 def forward(self, input: Tensor) -> Tensor:\r\n--> 114     return F.linear(input, self.weight, self.bias)\r\n\r\nRuntimeError: mat1 and mat2 shapes cannot be multiplied (1x4 and 768x650)\r\n```\r\n\r\n### Screenshots\r\n\r\n_No response_\r\n\r\n### Additional Context\r\n\r\n_No response_\r\n\r\n### Environment\r\n\r\n#### Versions:\r\n##### Flair\r\n0.11.3\r\n##### Pytorch\r\n1.13.0\r\n##### Transformers\r\n4.24.0\r\n#### GPU\r\nFalse\n", "before_files": [{"content": "import logging\nimport re\nfrom functools import lru_cache\nfrom pathlib import Path\nfrom typing import Any, Callable, Dict, List, Optional, Set, Union\nfrom unicodedata import category\n\nimport torch\n\nimport flair.embeddings\nimport flair.nn\nfrom flair.data import Dictionary, Sentence, Span\nfrom flair.file_utils import cached_path\n\nlog = logging.getLogger(\"flair\")\n\n\nclass CandidateGenerator:\n    \"\"\"\n    Given a string, the CandidateGenerator returns possible target classes as candidates.\n    \"\"\"\n\n    def __init__(self, candidates: Union[str, Dict], backoff: bool = True):\n        # internal candidate lists of generator\n        self.mention_to_candidates_map: Dict = {}\n\n        # load Zelda candidates if so passed\n        if isinstance(candidates, str) and candidates.lower() == \"zelda\":\n            zelda_path: str = \"https://flair.informatik.hu-berlin.de/resources/datasets/zelda\"\n            zelda_candidates = cached_path(f\"{zelda_path}/zelda_mention_entities_counter.pickle\", cache_dir=\"datasets\")\n            import pickle\n\n            with open(zelda_candidates, \"rb\") as handle:\n                mention_entities_counter = pickle.load(handle)\n\n            # create candidate lists\n            candidate_lists = {}\n            for mention in mention_entities_counter:\n                candidate_lists[mention] = list(mention_entities_counter[mention].keys())\n\n            self.mention_to_candidates_map = candidate_lists\n\n        elif isinstance(candidates, Dict):\n            self.mention_to_candidates_map = candidates\n\n        # if lower casing is enabled, create candidate lists of lower cased versions\n        self.backoff = backoff\n        if self.backoff:\n            # create a new dictionary for lower cased mentions\n            lowercased_mention_to_candidates_map: Dict = {}\n\n            # go through each mention and its candidates\n            for mention, candidates in self.mention_to_candidates_map.items():\n                backoff_mention = self._make_backoff_string(mention)\n                # check if backoff mention already seen. If so, add candidates. Else, create new entry.\n                if backoff_mention in lowercased_mention_to_candidates_map:\n                    current_candidates = lowercased_mention_to_candidates_map[backoff_mention]\n                    lowercased_mention_to_candidates_map[backoff_mention] = set(current_candidates).union(candidates)\n                else:\n                    lowercased_mention_to_candidates_map[backoff_mention] = candidates\n\n            # set lowercased version as map\n            self.mention_to_candidates_map = lowercased_mention_to_candidates_map\n\n    @lru_cache(maxsize=50000)\n    def _make_backoff_string(self, mention: str) -> str:\n        backoff_mention = mention.lower()\n        backoff_mention = \"\".join(ch for ch in backoff_mention if category(ch)[0] not in \"P\")\n        backoff_mention = re.sub(\" +\", \" \", backoff_mention)\n        return backoff_mention\n\n    def get_candidates(self, mention: str) -> Set[str]:\n        \"\"\"Given a mention, this method returns a set of candidate classes\"\"\"\n        if self.backoff:\n            mention = self._make_backoff_string(mention)\n\n        return set(self.mention_to_candidates_map[mention]) if mention in self.mention_to_candidates_map else set()\n\n\nclass EntityLinker(flair.nn.DefaultClassifier[Sentence, Span]):\n    \"\"\"\n    Entity Linking Model\n    The model expects text/sentences with annotated entity mentions and predicts entities to these mentions.\n    To this end a word embedding is used to embed the sentences and the embedding of the entity mention goes through a linear layer to get the actual class label.\n    The model is able to predict '<unk>' for entity mentions that the model can not confidently match to any of the known labels.\n    \"\"\"\n\n    def __init__(\n        self,\n        embeddings: flair.embeddings.TokenEmbeddings,\n        label_dictionary: Dictionary,\n        pooling_operation: str = \"first_last\",\n        label_type: str = \"nel\",\n        candidates: Optional[CandidateGenerator] = None,\n        **classifierargs,\n    ):\n        \"\"\"\n        Initializes an EntityLinker\n        :param embeddings: embeddings used to embed the words/sentences\n        :param label_dictionary: dictionary that gives ids to all classes. Should contain <unk>\n        :param pooling_operation: either 'average', 'first', 'last' or 'first&last'. Specifies the way of how text representations of entity mentions (with more than one word) are handled.\n        E.g. 'average' means that as text representation we take the average of the embeddings of the words in the mention. 'first&last' concatenates\n        the embedding of the first and the embedding of the last word.\n        :param label_type: name of the label you use.\n        \"\"\"\n\n        super(EntityLinker, self).__init__(\n            embeddings=embeddings,\n            label_dictionary=label_dictionary,\n            final_embedding_size=embeddings.embedding_length * 2\n            if pooling_operation == \"first_last\"\n            else embeddings.embedding_length,\n            **classifierargs,\n        )\n\n        self.pooling_operation = pooling_operation\n        self._label_type = label_type\n\n        cases: Dict[str, Callable[[Span, List[str]], torch.Tensor]] = {\n            \"average\": self.emb_mean,\n            \"first\": self.emb_first,\n            \"last\": self.emb_last,\n            \"first_last\": self.emb_firstAndLast,\n        }\n\n        if pooling_operation not in cases:\n            raise KeyError('pooling_operation has to be one of \"average\", \"first\", \"last\" or \"first_last\"')\n\n        self.aggregated_embedding = cases[pooling_operation]\n\n        self.candidates = candidates\n\n        self.to(flair.device)\n\n    def emb_first(self, span: Span, embedding_names):\n        return span.tokens[0].get_embedding(embedding_names)\n\n    def emb_last(self, span: Span, embedding_names):\n        return span.tokens[-1].get_embedding(embedding_names)\n\n    def emb_firstAndLast(self, span: Span, embedding_names):\n        return torch.cat(\n            (span.tokens[0].get_embedding(embedding_names), span.tokens[-1].get_embedding(embedding_names)), 0\n        )\n\n    def emb_mean(self, span, embedding_names):\n        return torch.mean(torch.cat([token.get_embedding(embedding_names) for token in span], 0), 0)\n\n    def _get_data_points_from_sentence(self, sentence: Sentence) -> List[Span]:\n        return sentence.get_spans(self.label_type)\n\n    def _filter_data_point(self, data_point: Sentence) -> bool:\n        return bool(data_point.get_labels(self.label_type))\n\n    def _get_embedding_for_data_point(self, prediction_data_point: Span) -> torch.Tensor:\n        return self.aggregated_embedding(prediction_data_point, self.embeddings.get_names())\n\n    def _get_state_dict(self):\n        model_state = {\n            **super()._get_state_dict(),\n            \"word_embeddings\": self.embeddings.save_embeddings(use_state_dict=False),\n            \"label_type\": self.label_type,\n            \"label_dictionary\": self.label_dictionary,\n            \"pooling_operation\": self.pooling_operation,\n            \"loss_weights\": self.weight_dict,\n            \"candidates\": self.candidates,\n        }\n        return model_state\n\n    def _print_predictions(self, batch, gold_label_type):\n        lines = []\n        for datapoint in batch:\n            eval_line = f\"\\n{datapoint.to_original_text()}\\n\"\n\n            for span in datapoint.get_spans(gold_label_type):\n                symbol = \"\u2713\" if span.get_label(gold_label_type).value == span.get_label(\"predicted\").value else \"\u274c\"\n                eval_line += (\n                    f' - \"{span.text}\" / {span.get_label(gold_label_type).value}'\n                    f' --> {span.get_label(\"predicted\").value} ({symbol})\\n'\n                )\n\n            lines.append(eval_line)\n        return lines\n\n    @classmethod\n    def _init_model_with_state_dict(cls, state, **kwargs):\n        # remap state dict for models serialized with Flair <= 0.11.3\n        import re\n\n        state_dict = state[\"state_dict\"]\n        for key in list(state_dict.keys()):\n            state_dict[re.sub(\"^word_embeddings\\\\.\", \"embeddings.\", key)] = state_dict.pop(key)\n\n        return super()._init_model_with_state_dict(\n            state,\n            embeddings=state.get(\"word_embeddings\"),\n            label_dictionary=state.get(\"label_dictionary\"),\n            label_type=state.get(\"label_type\"),\n            pooling_operation=state.get(\"pooling_operation\"),\n            loss_weights=state.get(\"loss_weights\", {\"<unk>\": 0.3}),\n            candidates=state.get(\"candidates\", None),\n            **kwargs,\n        )\n\n    @property\n    def label_type(self):\n        return self._label_type\n\n    def _mask_scores(self, scores: torch.Tensor, data_points: List[Span]):\n        if not self.candidates:\n            return scores\n\n        masked_scores = -torch.inf * torch.ones(scores.size(), requires_grad=True, device=flair.device)\n\n        for idx, span in enumerate(data_points):\n            # get the candidates\n            candidate_set = self.candidates.get_candidates(span.text)\n            # during training, add the gold value as candidate\n            if self.training:\n                candidate_set.add(span.get_label(self.label_type).value)\n            candidate_set.add(\"<unk>\")\n            indices_of_candidates = [self.label_dictionary.get_idx_for_item(candidate) for candidate in candidate_set]\n            masked_scores[idx, indices_of_candidates] = scores[idx, indices_of_candidates]\n\n        return masked_scores\n\n    @classmethod\n    def load(cls, model_path: Union[str, Path, Dict[str, Any]]) -> \"EntityLinker\":\n        from typing import cast\n\n        return cast(\"EntityLinker\", super().load(model_path=model_path))\n", "path": "flair/models/entity_linker_model.py"}], "after_files": [{"content": "import logging\nimport re\nfrom functools import lru_cache\nfrom pathlib import Path\nfrom typing import Any, Callable, Dict, List, Optional, Set, Union\nfrom unicodedata import category\n\nimport torch\n\nimport flair.embeddings\nimport flair.nn\nfrom flair.data import Dictionary, Sentence, Span\nfrom flair.file_utils import cached_path\n\nlog = logging.getLogger(\"flair\")\n\n\nclass CandidateGenerator:\n    \"\"\"\n    Given a string, the CandidateGenerator returns possible target classes as candidates.\n    \"\"\"\n\n    def __init__(self, candidates: Union[str, Dict], backoff: bool = True):\n        # internal candidate lists of generator\n        self.mention_to_candidates_map: Dict = {}\n\n        # load Zelda candidates if so passed\n        if isinstance(candidates, str) and candidates.lower() == \"zelda\":\n            zelda_path: str = \"https://flair.informatik.hu-berlin.de/resources/datasets/zelda\"\n            zelda_candidates = cached_path(f\"{zelda_path}/zelda_mention_entities_counter.pickle\", cache_dir=\"datasets\")\n            import pickle\n\n            with open(zelda_candidates, \"rb\") as handle:\n                mention_entities_counter = pickle.load(handle)\n\n            # create candidate lists\n            candidate_lists = {}\n            for mention in mention_entities_counter:\n                candidate_lists[mention] = list(mention_entities_counter[mention].keys())\n\n            self.mention_to_candidates_map = candidate_lists\n\n        elif isinstance(candidates, Dict):\n            self.mention_to_candidates_map = candidates\n\n        # if lower casing is enabled, create candidate lists of lower cased versions\n        self.backoff = backoff\n        if self.backoff:\n            # create a new dictionary for lower cased mentions\n            lowercased_mention_to_candidates_map: Dict = {}\n\n            # go through each mention and its candidates\n            for mention, candidates in self.mention_to_candidates_map.items():\n                backoff_mention = self._make_backoff_string(mention)\n                # check if backoff mention already seen. If so, add candidates. Else, create new entry.\n                if backoff_mention in lowercased_mention_to_candidates_map:\n                    current_candidates = lowercased_mention_to_candidates_map[backoff_mention]\n                    lowercased_mention_to_candidates_map[backoff_mention] = set(current_candidates).union(candidates)\n                else:\n                    lowercased_mention_to_candidates_map[backoff_mention] = candidates\n\n            # set lowercased version as map\n            self.mention_to_candidates_map = lowercased_mention_to_candidates_map\n\n    @lru_cache(maxsize=50000)\n    def _make_backoff_string(self, mention: str) -> str:\n        backoff_mention = mention.lower()\n        backoff_mention = \"\".join(ch for ch in backoff_mention if category(ch)[0] not in \"P\")\n        backoff_mention = re.sub(\" +\", \" \", backoff_mention)\n        return backoff_mention\n\n    def get_candidates(self, mention: str) -> Set[str]:\n        \"\"\"Given a mention, this method returns a set of candidate classes\"\"\"\n        if self.backoff:\n            mention = self._make_backoff_string(mention)\n\n        return set(self.mention_to_candidates_map[mention]) if mention in self.mention_to_candidates_map else set()\n\n\nclass EntityLinker(flair.nn.DefaultClassifier[Sentence, Span]):\n    \"\"\"\n    Entity Linking Model\n    The model expects text/sentences with annotated entity mentions and predicts entities to these mentions.\n    To this end a word embedding is used to embed the sentences and the embedding of the entity mention goes through a linear layer to get the actual class label.\n    The model is able to predict '<unk>' for entity mentions that the model can not confidently match to any of the known labels.\n    \"\"\"\n\n    def __init__(\n        self,\n        embeddings: flair.embeddings.TokenEmbeddings,\n        label_dictionary: Dictionary,\n        pooling_operation: str = \"first_last\",\n        label_type: str = \"nel\",\n        candidates: Optional[CandidateGenerator] = None,\n        **classifierargs,\n    ):\n        \"\"\"\n        Initializes an EntityLinker\n        :param embeddings: embeddings used to embed the words/sentences\n        :param label_dictionary: dictionary that gives ids to all classes. Should contain <unk>\n        :param pooling_operation: either 'average', 'first', 'last' or 'first&last'. Specifies the way of how text representations of entity mentions (with more than one word) are handled.\n        E.g. 'average' means that as text representation we take the average of the embeddings of the words in the mention. 'first&last' concatenates\n        the embedding of the first and the embedding of the last word.\n        :param label_type: name of the label you use.\n        \"\"\"\n\n        super(EntityLinker, self).__init__(\n            embeddings=embeddings,\n            label_dictionary=label_dictionary,\n            final_embedding_size=embeddings.embedding_length * 2\n            if pooling_operation == \"first_last\"\n            else embeddings.embedding_length,\n            **classifierargs,\n        )\n\n        self.pooling_operation = pooling_operation\n        self._label_type = label_type\n\n        cases: Dict[str, Callable[[Span, List[str]], torch.Tensor]] = {\n            \"average\": self.emb_mean,\n            \"first\": self.emb_first,\n            \"last\": self.emb_last,\n            \"first_last\": self.emb_firstAndLast,\n        }\n\n        if pooling_operation not in cases:\n            raise KeyError('pooling_operation has to be one of \"average\", \"first\", \"last\" or \"first_last\"')\n\n        self.aggregated_embedding = cases[pooling_operation]\n\n        self.candidates = candidates\n\n        self.to(flair.device)\n\n    def emb_first(self, span: Span, embedding_names):\n        return span.tokens[0].get_embedding(embedding_names)\n\n    def emb_last(self, span: Span, embedding_names):\n        return span.tokens[-1].get_embedding(embedding_names)\n\n    def emb_firstAndLast(self, span: Span, embedding_names):\n        return torch.cat(\n            (span.tokens[0].get_embedding(embedding_names), span.tokens[-1].get_embedding(embedding_names)), 0\n        )\n\n    def emb_mean(self, span, embedding_names):\n        return torch.mean(torch.stack([token.get_embedding(embedding_names) for token in span], 0), 0)\n\n    def _get_data_points_from_sentence(self, sentence: Sentence) -> List[Span]:\n        return sentence.get_spans(self.label_type)\n\n    def _filter_data_point(self, data_point: Sentence) -> bool:\n        return bool(data_point.get_labels(self.label_type))\n\n    def _get_embedding_for_data_point(self, prediction_data_point: Span) -> torch.Tensor:\n        return self.aggregated_embedding(prediction_data_point, self.embeddings.get_names())\n\n    def _get_state_dict(self):\n        model_state = {\n            **super()._get_state_dict(),\n            \"word_embeddings\": self.embeddings.save_embeddings(use_state_dict=False),\n            \"label_type\": self.label_type,\n            \"label_dictionary\": self.label_dictionary,\n            \"pooling_operation\": self.pooling_operation,\n            \"loss_weights\": self.weight_dict,\n            \"candidates\": self.candidates,\n        }\n        return model_state\n\n    def _print_predictions(self, batch, gold_label_type):\n        lines = []\n        for datapoint in batch:\n            eval_line = f\"\\n{datapoint.to_original_text()}\\n\"\n\n            for span in datapoint.get_spans(gold_label_type):\n                symbol = \"\u2713\" if span.get_label(gold_label_type).value == span.get_label(\"predicted\").value else \"\u274c\"\n                eval_line += (\n                    f' - \"{span.text}\" / {span.get_label(gold_label_type).value}'\n                    f' --> {span.get_label(\"predicted\").value} ({symbol})\\n'\n                )\n\n            lines.append(eval_line)\n        return lines\n\n    @classmethod\n    def _init_model_with_state_dict(cls, state, **kwargs):\n        # remap state dict for models serialized with Flair <= 0.11.3\n        import re\n\n        state_dict = state[\"state_dict\"]\n        for key in list(state_dict.keys()):\n            state_dict[re.sub(\"^word_embeddings\\\\.\", \"embeddings.\", key)] = state_dict.pop(key)\n\n        return super()._init_model_with_state_dict(\n            state,\n            embeddings=state.get(\"word_embeddings\"),\n            label_dictionary=state.get(\"label_dictionary\"),\n            label_type=state.get(\"label_type\"),\n            pooling_operation=state.get(\"pooling_operation\"),\n            loss_weights=state.get(\"loss_weights\", {\"<unk>\": 0.3}),\n            candidates=state.get(\"candidates\", None),\n            **kwargs,\n        )\n\n    @property\n    def label_type(self):\n        return self._label_type\n\n    def _mask_scores(self, scores: torch.Tensor, data_points: List[Span]):\n        if not self.candidates:\n            return scores\n\n        masked_scores = -torch.inf * torch.ones(scores.size(), requires_grad=True, device=flair.device)\n\n        for idx, span in enumerate(data_points):\n            # get the candidates\n            candidate_set = self.candidates.get_candidates(span.text)\n            # during training, add the gold value as candidate\n            if self.training:\n                candidate_set.add(span.get_label(self.label_type).value)\n            candidate_set.add(\"<unk>\")\n            indices_of_candidates = [self.label_dictionary.get_idx_for_item(candidate) for candidate in candidate_set]\n            masked_scores[idx, indices_of_candidates] = scores[idx, indices_of_candidates]\n\n        return masked_scores\n\n    @classmethod\n    def load(cls, model_path: Union[str, Path, Dict[str, Any]]) -> \"EntityLinker\":\n        from typing import cast\n\n        return cast(\"EntityLinker\", super().load(model_path=model_path))\n", "path": "flair/models/entity_linker_model.py"}]}