relik-ie commited on
Commit
d844ec0
1 Parent(s): 84ec6d1

Upload model

Browse files
Files changed (5) hide show
  1. README.md +199 -0
  2. config.json +26 -0
  3. configuration_relik.py +45 -0
  4. model.safetensors +3 -0
  5. modeling_relik.py +999 -0
README.md ADDED
@@ -0,0 +1,199 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ library_name: transformers
3
+ tags: []
4
+ ---
5
+
6
+ # Model Card for Model ID
7
+
8
+ <!-- Provide a quick summary of what the model is/does. -->
9
+
10
+
11
+
12
+ ## Model Details
13
+
14
+ ### Model Description
15
+
16
+ <!-- Provide a longer summary of what this model is. -->
17
+
18
+ This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
19
+
20
+ - **Developed by:** [More Information Needed]
21
+ - **Funded by [optional]:** [More Information Needed]
22
+ - **Shared by [optional]:** [More Information Needed]
23
+ - **Model type:** [More Information Needed]
24
+ - **Language(s) (NLP):** [More Information Needed]
25
+ - **License:** [More Information Needed]
26
+ - **Finetuned from model [optional]:** [More Information Needed]
27
+
28
+ ### Model Sources [optional]
29
+
30
+ <!-- Provide the basic links for the model. -->
31
+
32
+ - **Repository:** [More Information Needed]
33
+ - **Paper [optional]:** [More Information Needed]
34
+ - **Demo [optional]:** [More Information Needed]
35
+
36
+ ## Uses
37
+
38
+ <!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
39
+
40
+ ### Direct Use
41
+
42
+ <!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
43
+
44
+ [More Information Needed]
45
+
46
+ ### Downstream Use [optional]
47
+
48
+ <!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
49
+
50
+ [More Information Needed]
51
+
52
+ ### Out-of-Scope Use
53
+
54
+ <!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
55
+
56
+ [More Information Needed]
57
+
58
+ ## Bias, Risks, and Limitations
59
+
60
+ <!-- This section is meant to convey both technical and sociotechnical limitations. -->
61
+
62
+ [More Information Needed]
63
+
64
+ ### Recommendations
65
+
66
+ <!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
67
+
68
+ Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
69
+
70
+ ## How to Get Started with the Model
71
+
72
+ Use the code below to get started with the model.
73
+
74
+ [More Information Needed]
75
+
76
+ ## Training Details
77
+
78
+ ### Training Data
79
+
80
+ <!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
81
+
82
+ [More Information Needed]
83
+
84
+ ### Training Procedure
85
+
86
+ <!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
87
+
88
+ #### Preprocessing [optional]
89
+
90
+ [More Information Needed]
91
+
92
+
93
+ #### Training Hyperparameters
94
+
95
+ - **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
96
+
97
+ #### Speeds, Sizes, Times [optional]
98
+
99
+ <!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
100
+
101
+ [More Information Needed]
102
+
103
+ ## Evaluation
104
+
105
+ <!-- This section describes the evaluation protocols and provides the results. -->
106
+
107
+ ### Testing Data, Factors & Metrics
108
+
109
+ #### Testing Data
110
+
111
+ <!-- This should link to a Dataset Card if possible. -->
112
+
113
+ [More Information Needed]
114
+
115
+ #### Factors
116
+
117
+ <!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
118
+
119
+ [More Information Needed]
120
+
121
+ #### Metrics
122
+
123
+ <!-- These are the evaluation metrics being used, ideally with a description of why. -->
124
+
125
+ [More Information Needed]
126
+
127
+ ### Results
128
+
129
+ [More Information Needed]
130
+
131
+ #### Summary
132
+
133
+
134
+
135
+ ## Model Examination [optional]
136
+
137
+ <!-- Relevant interpretability work for the model goes here -->
138
+
139
+ [More Information Needed]
140
+
141
+ ## Environmental Impact
142
+
143
+ <!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
144
+
145
+ Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
146
+
147
+ - **Hardware Type:** [More Information Needed]
148
+ - **Hours used:** [More Information Needed]
149
+ - **Cloud Provider:** [More Information Needed]
150
+ - **Compute Region:** [More Information Needed]
151
+ - **Carbon Emitted:** [More Information Needed]
152
+
153
+ ## Technical Specifications [optional]
154
+
155
+ ### Model Architecture and Objective
156
+
157
+ [More Information Needed]
158
+
159
+ ### Compute Infrastructure
160
+
161
+ [More Information Needed]
162
+
163
+ #### Hardware
164
+
165
+ [More Information Needed]
166
+
167
+ #### Software
168
+
169
+ [More Information Needed]
170
+
171
+ ## Citation [optional]
172
+
173
+ <!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
174
+
175
+ **BibTeX:**
176
+
177
+ [More Information Needed]
178
+
179
+ **APA:**
180
+
181
+ [More Information Needed]
182
+
183
+ ## Glossary [optional]
184
+
185
+ <!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
186
+
187
+ [More Information Needed]
188
+
189
+ ## More Information [optional]
190
+
191
+ [More Information Needed]
192
+
193
+ ## Model Card Authors [optional]
194
+
195
+ [More Information Needed]
196
+
197
+ ## Model Card Contact
198
+
199
+ [More Information Needed]
config.json ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "_name_or_path": "/media/ssd/perelluis/relik_experiments/models/relik-reader-deberta-v3-large-re-wikipedia",
3
+ "activation": "gelu",
4
+ "add_entity_embedding": null,
5
+ "additional_special_symbols": 30,
6
+ "additional_special_symbols_types": 0,
7
+ "architectures": [
8
+ "RelikReaderREModel"
9
+ ],
10
+ "auto_map": {
11
+ "AutoConfig": "configuration_relik.RelikReaderConfig",
12
+ "AutoModel": "modeling_relik.RelikReaderREModel"
13
+ },
14
+ "binary_end_logits": false,
15
+ "default_reader_class": null,
16
+ "entity_type_loss": null,
17
+ "linears_hidden_size": 512,
18
+ "model_type": "relik-reader",
19
+ "num_layers": null,
20
+ "threshold": 0.5,
21
+ "torch_dtype": "float32",
22
+ "training": true,
23
+ "transformer_model": "microsoft/deberta-v3-small",
24
+ "transformers_version": "4.41.2",
25
+ "use_last_k_layers": 1
26
+ }
configuration_relik.py ADDED
@@ -0,0 +1,45 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Optional
2
+
3
+ from transformers import AutoConfig
4
+ from transformers.configuration_utils import PretrainedConfig
5
+
6
+
7
+ class RelikReaderConfig(PretrainedConfig):
8
+ model_type = "relik-reader"
9
+
10
+ def __init__(
11
+ self,
12
+ transformer_model: str = "microsoft/deberta-v3-base",
13
+ additional_special_symbols: int = 101,
14
+ additional_special_symbols_types: Optional[int] = 0,
15
+ num_layers: Optional[int] = None,
16
+ activation: str = "gelu",
17
+ linears_hidden_size: Optional[int] = 512,
18
+ use_last_k_layers: int = 1,
19
+ entity_type_loss: bool = False,
20
+ add_entity_embedding: bool = None,
21
+ binary_end_logits: bool = False,
22
+ training: bool = False,
23
+ default_reader_class: Optional[str] = None,
24
+ threshold: Optional[float] = 0.5,
25
+ **kwargs
26
+ ) -> None:
27
+ # TODO: add name_or_path to kwargs
28
+ self.transformer_model = transformer_model
29
+ self.additional_special_symbols = additional_special_symbols
30
+ self.additional_special_symbols_types = additional_special_symbols_types
31
+ self.num_layers = num_layers
32
+ self.activation = activation
33
+ self.linears_hidden_size = linears_hidden_size
34
+ self.use_last_k_layers = use_last_k_layers
35
+ self.entity_type_loss = entity_type_loss
36
+ self.add_entity_embedding = (
37
+ True
38
+ if add_entity_embedding is None and entity_type_loss
39
+ else add_entity_embedding
40
+ )
41
+ self.threshold = threshold
42
+ self.binary_end_logits = binary_end_logits
43
+ self.training = training
44
+ self.default_reader_class = default_reader_class
45
+ super().__init__(**kwargs)
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fcaad405b58a947631c5690834dad6d3dcaf1475446a90cbfe6ec5523a44eb8e
3
+ size 585824816
modeling_relik.py ADDED
@@ -0,0 +1,999 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, Dict, Optional
2
+
3
+ import torch
4
+ from transformers import AutoModel, PreTrainedModel
5
+ from transformers.activations import ClippedGELUActivation, GELUActivation
6
+ from transformers.configuration_utils import PretrainedConfig
7
+ from transformers.modeling_utils import PoolerEndLogits
8
+
9
+ from .configuration_relik import RelikReaderConfig
10
+
11
+
12
+ class RelikReaderSample:
13
+ def __init__(self, **kwargs):
14
+ super().__setattr__("_d", {})
15
+ self._d = kwargs
16
+
17
+ def __getattribute__(self, item):
18
+ return super(RelikReaderSample, self).__getattribute__(item)
19
+
20
+ def __getattr__(self, item):
21
+ if item.startswith("__") and item.endswith("__"):
22
+ # this is likely some python library-specific variable (such as __deepcopy__ for copy)
23
+ # better follow standard behavior here
24
+ raise AttributeError(item)
25
+ elif item in self._d:
26
+ return self._d[item]
27
+ else:
28
+ return None
29
+
30
+ def __setattr__(self, key, value):
31
+ if key in self._d:
32
+ self._d[key] = value
33
+ else:
34
+ super().__setattr__(key, value)
35
+ self._d[key] = value
36
+
37
+
38
+ activation2functions = {
39
+ "relu": torch.nn.ReLU(),
40
+ "gelu": GELUActivation(),
41
+ "gelu_10": ClippedGELUActivation(-10, 10),
42
+ }
43
+
44
+
45
+ class PoolerEndLogitsBi(PoolerEndLogits):
46
+ def __init__(self, config: PretrainedConfig):
47
+ super().__init__(config)
48
+ self.dense_1 = torch.nn.Linear(config.hidden_size, 2)
49
+
50
+ def forward(
51
+ self,
52
+ hidden_states: torch.FloatTensor,
53
+ start_states: Optional[torch.FloatTensor] = None,
54
+ start_positions: Optional[torch.LongTensor] = None,
55
+ p_mask: Optional[torch.FloatTensor] = None,
56
+ ) -> torch.FloatTensor:
57
+ if p_mask is not None:
58
+ p_mask = p_mask.unsqueeze(-1)
59
+ logits = super().forward(
60
+ hidden_states,
61
+ start_states,
62
+ start_positions,
63
+ p_mask,
64
+ )
65
+ return logits
66
+
67
+
68
+ class RelikReaderSpanModel(PreTrainedModel):
69
+ config_class = RelikReaderConfig
70
+
71
+ def __init__(self, config: RelikReaderConfig, *args, **kwargs):
72
+ super().__init__(config)
73
+ # Transformer model declaration
74
+ self.config = config
75
+ self.transformer_model = (
76
+ AutoModel.from_pretrained(self.config.transformer_model)
77
+ if self.config.num_layers is None
78
+ else AutoModel.from_pretrained(
79
+ self.config.transformer_model, num_hidden_layers=self.config.num_layers
80
+ )
81
+ )
82
+ self.transformer_model.resize_token_embeddings(
83
+ self.transformer_model.config.vocab_size
84
+ + self.config.additional_special_symbols
85
+ )
86
+
87
+ self.activation = self.config.activation
88
+ self.linears_hidden_size = self.config.linears_hidden_size
89
+ self.use_last_k_layers = self.config.use_last_k_layers
90
+
91
+ # named entity detection layers
92
+ self.ned_start_classifier = self._get_projection_layer(
93
+ self.activation, last_hidden=2, layer_norm=False
94
+ )
95
+ if self.config.binary_end_logits:
96
+ self.ned_end_classifier = PoolerEndLogitsBi(self.transformer_model.config)
97
+ else:
98
+ self.ned_end_classifier = PoolerEndLogits(self.transformer_model.config)
99
+
100
+ # END entity disambiguation layer
101
+ self.ed_start_projector = self._get_projection_layer(self.activation)
102
+ self.ed_end_projector = self._get_projection_layer(self.activation)
103
+
104
+ self.training = self.config.training
105
+
106
+ # criterion
107
+ self.criterion = torch.nn.CrossEntropyLoss()
108
+
109
+ def _get_projection_layer(
110
+ self,
111
+ activation: str,
112
+ last_hidden: Optional[int] = None,
113
+ input_hidden=None,
114
+ layer_norm: bool = True,
115
+ ) -> torch.nn.Sequential:
116
+ head_components = [
117
+ torch.nn.Dropout(0.1),
118
+ torch.nn.Linear(
119
+ (
120
+ self.transformer_model.config.hidden_size * self.use_last_k_layers
121
+ if input_hidden is None
122
+ else input_hidden
123
+ ),
124
+ self.linears_hidden_size,
125
+ ),
126
+ activation2functions[activation],
127
+ torch.nn.Dropout(0.1),
128
+ torch.nn.Linear(
129
+ self.linears_hidden_size,
130
+ self.linears_hidden_size if last_hidden is None else last_hidden,
131
+ ),
132
+ ]
133
+
134
+ if layer_norm:
135
+ head_components.append(
136
+ torch.nn.LayerNorm(
137
+ self.linears_hidden_size if last_hidden is None else last_hidden,
138
+ self.transformer_model.config.layer_norm_eps,
139
+ )
140
+ )
141
+
142
+ return torch.nn.Sequential(*head_components)
143
+
144
+ def _mask_logits(self, logits: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
145
+ mask = mask.unsqueeze(-1)
146
+ if next(self.parameters()).dtype == torch.float16:
147
+ logits = logits * (1 - mask) - 65500 * mask
148
+ else:
149
+ logits = logits * (1 - mask) - 1e30 * mask
150
+ return logits
151
+
152
+ def _get_model_features(
153
+ self,
154
+ input_ids: torch.Tensor,
155
+ attention_mask: torch.Tensor,
156
+ token_type_ids: Optional[torch.Tensor],
157
+ ):
158
+ model_input = {
159
+ "input_ids": input_ids,
160
+ "attention_mask": attention_mask,
161
+ "output_hidden_states": self.use_last_k_layers > 1,
162
+ }
163
+
164
+ if token_type_ids is not None:
165
+ model_input["token_type_ids"] = token_type_ids
166
+
167
+ model_output = self.transformer_model(**model_input)
168
+
169
+ if self.use_last_k_layers > 1:
170
+ model_features = torch.cat(
171
+ model_output[1][-self.use_last_k_layers :], dim=-1
172
+ )
173
+ else:
174
+ model_features = model_output[0]
175
+
176
+ return model_features
177
+
178
+ def compute_ned_end_logits(
179
+ self,
180
+ start_predictions,
181
+ start_labels,
182
+ model_features,
183
+ prediction_mask,
184
+ batch_size,
185
+ ) -> Optional[torch.Tensor]:
186
+ # todo: maybe when constraining on the spans,
187
+ # we should not use a prediction_mask for the end tokens.
188
+ # at least we should not during training imo
189
+ start_positions = start_labels if self.training else start_predictions
190
+ start_positions_indices = (
191
+ torch.arange(start_positions.size(1), device=start_positions.device)
192
+ .unsqueeze(0)
193
+ .expand(batch_size, -1)[start_positions > 0]
194
+ ).to(start_positions.device)
195
+
196
+ if len(start_positions_indices) > 0:
197
+ expanded_features = model_features.repeat_interleave(
198
+ torch.sum(start_positions > 0, dim=-1), dim=0
199
+ )
200
+ expanded_prediction_mask = prediction_mask.repeat_interleave(
201
+ torch.sum(start_positions > 0, dim=-1), dim=0
202
+ )
203
+ end_logits = self.ned_end_classifier(
204
+ hidden_states=expanded_features,
205
+ start_positions=start_positions_indices,
206
+ p_mask=expanded_prediction_mask,
207
+ )
208
+
209
+ return end_logits
210
+
211
+ return None
212
+
213
+ def compute_classification_logits(
214
+ self,
215
+ model_features_start,
216
+ model_features_end,
217
+ special_symbols_features,
218
+ ) -> torch.Tensor:
219
+ model_start_features = self.ed_start_projector(model_features_start)
220
+ model_end_features = self.ed_end_projector(model_features_end)
221
+ model_start_features_symbols = self.ed_start_projector(special_symbols_features)
222
+ model_end_features_symbols = self.ed_end_projector(special_symbols_features)
223
+
224
+ model_ed_features = torch.cat(
225
+ [model_start_features, model_end_features], dim=-1
226
+ )
227
+ special_symbols_representation = torch.cat(
228
+ [model_start_features_symbols, model_end_features_symbols], dim=-1
229
+ )
230
+
231
+ logits = torch.bmm(
232
+ model_ed_features,
233
+ torch.permute(special_symbols_representation, (0, 2, 1)),
234
+ )
235
+
236
+ logits = self._mask_logits(
237
+ logits, (model_features_start == -100).all(2).long()
238
+ )
239
+ return logits
240
+
241
+ def forward(
242
+ self,
243
+ input_ids: torch.Tensor,
244
+ attention_mask: torch.Tensor,
245
+ token_type_ids: Optional[torch.Tensor] = None,
246
+ prediction_mask: Optional[torch.Tensor] = None,
247
+ special_symbols_mask: Optional[torch.Tensor] = None,
248
+ start_labels: Optional[torch.Tensor] = None,
249
+ end_labels: Optional[torch.Tensor] = None,
250
+ use_predefined_spans: bool = False,
251
+ *args,
252
+ **kwargs,
253
+ ) -> Dict[str, Any]:
254
+ batch_size, seq_len = input_ids.shape
255
+
256
+ model_features = self._get_model_features(
257
+ input_ids, attention_mask, token_type_ids
258
+ )
259
+
260
+ ned_start_labels = None
261
+
262
+ # named entity detection if required
263
+ if use_predefined_spans: # no need to compute spans
264
+ ned_start_logits, ned_start_probabilities, ned_start_predictions = (
265
+ None,
266
+ None,
267
+ (
268
+ torch.clone(start_labels)
269
+ if start_labels is not None
270
+ else torch.zeros_like(input_ids)
271
+ ),
272
+ )
273
+ ned_end_logits, ned_end_probabilities, ned_end_predictions = (
274
+ None,
275
+ None,
276
+ (
277
+ torch.clone(end_labels)
278
+ if end_labels is not None
279
+ else torch.zeros_like(input_ids)
280
+ ),
281
+ )
282
+ ned_start_predictions[ned_start_predictions > 0] = 1
283
+ ned_end_predictions[end_labels > 0] = 1
284
+ ned_end_predictions = ned_end_predictions[~(end_labels == -100).all(2)]
285
+
286
+ else: # compute spans
287
+ # start boundary prediction
288
+ ned_start_logits = self.ned_start_classifier(model_features)
289
+ ned_start_logits = self._mask_logits(ned_start_logits, prediction_mask)
290
+ ned_start_probabilities = torch.softmax(ned_start_logits, dim=-1)
291
+ ned_start_predictions = ned_start_probabilities.argmax(dim=-1)
292
+
293
+ # end boundary prediction
294
+ ned_start_labels = (
295
+ torch.zeros_like(start_labels) if start_labels is not None else None
296
+ )
297
+
298
+ if ned_start_labels is not None:
299
+ ned_start_labels[start_labels == -100] = -100
300
+ ned_start_labels[start_labels > 0] = 1
301
+
302
+ ned_end_logits = self.compute_ned_end_logits(
303
+ ned_start_predictions,
304
+ ned_start_labels,
305
+ model_features,
306
+ prediction_mask,
307
+ batch_size,
308
+ )
309
+
310
+ if ned_end_logits is not None:
311
+ ned_end_probabilities = torch.softmax(ned_end_logits, dim=-1)
312
+ if not self.config.binary_end_logits:
313
+ ned_end_predictions = torch.argmax(ned_end_probabilities, dim=-1, keepdim=True)
314
+ ned_end_predictions = torch.zeros_like(ned_end_probabilities).scatter_(1, ned_end_predictions, 1)
315
+ else:
316
+ ned_end_predictions = torch.argmax(ned_end_probabilities, dim=-1)
317
+ else:
318
+ ned_end_logits, ned_end_probabilities = None, None
319
+ ned_end_predictions = ned_start_predictions.new_zeros(batch_size, seq_len)
320
+
321
+ if not self.training:
322
+ # if len(ned_end_predictions.shape) < 2:
323
+ # print(ned_end_predictions)
324
+ end_preds_count = ned_end_predictions.sum(1)
325
+ # If there are no end predictions for a start prediction, remove the start prediction
326
+ if (end_preds_count == 0).any() and (ned_start_predictions > 0).any():
327
+ ned_start_predictions[ned_start_predictions == 1] = (
328
+ end_preds_count != 0
329
+ ).long()
330
+ ned_end_predictions = ned_end_predictions[end_preds_count != 0]
331
+
332
+ if end_labels is not None:
333
+ end_labels = end_labels[~(end_labels == -100).all(2)]
334
+
335
+ start_position, end_position = (
336
+ (start_labels, end_labels)
337
+ if self.training
338
+ else (ned_start_predictions, ned_end_predictions)
339
+ )
340
+ start_counts = (start_position > 0).sum(1)
341
+ if (start_counts > 0).any():
342
+ ned_end_predictions = ned_end_predictions.split(start_counts.tolist())
343
+ # Entity disambiguation
344
+ if (end_position > 0).sum() > 0:
345
+ ends_count = (end_position > 0).sum(1)
346
+ model_entity_start = torch.repeat_interleave(
347
+ model_features[start_position > 0], ends_count, dim=0
348
+ )
349
+ model_entity_end = torch.repeat_interleave(
350
+ model_features, start_counts, dim=0)[
351
+ end_position > 0
352
+ ]
353
+ ents_count = torch.nn.utils.rnn.pad_sequence(
354
+ torch.split(ends_count, start_counts.tolist()),
355
+ batch_first=True,
356
+ padding_value=0,
357
+ ).sum(1)
358
+
359
+ model_entity_start = torch.nn.utils.rnn.pad_sequence(
360
+ torch.split(model_entity_start, ents_count.tolist()),
361
+ batch_first=True,
362
+ padding_value=-100,
363
+ )
364
+
365
+ model_entity_end = torch.nn.utils.rnn.pad_sequence(
366
+ torch.split(model_entity_end, ents_count.tolist()),
367
+ batch_first=True,
368
+ padding_value=-100,
369
+ )
370
+
371
+ ed_logits = self.compute_classification_logits(
372
+ model_entity_start,
373
+ model_entity_end,
374
+ model_features[special_symbols_mask].view(
375
+ batch_size, -1, model_features.shape[-1]
376
+ ),
377
+ )
378
+ ed_probabilities = torch.softmax(ed_logits, dim=-1)
379
+ ed_predictions = torch.argmax(ed_probabilities, dim=-1)
380
+ else:
381
+ ed_logits, ed_probabilities, ed_predictions = (
382
+ None,
383
+ ned_start_predictions.new_zeros(batch_size, seq_len),
384
+ ned_start_predictions.new_zeros(batch_size),
385
+ )
386
+ # output build
387
+ output_dict = dict(
388
+ batch_size=batch_size,
389
+ ned_start_logits=ned_start_logits,
390
+ ned_start_probabilities=ned_start_probabilities,
391
+ ned_start_predictions=ned_start_predictions,
392
+ ned_end_logits=ned_end_logits,
393
+ ned_end_probabilities=ned_end_probabilities,
394
+ ned_end_predictions=ned_end_predictions,
395
+ ed_logits=ed_logits,
396
+ ed_probabilities=ed_probabilities,
397
+ ed_predictions=ed_predictions,
398
+ )
399
+
400
+ # compute loss if labels
401
+ if start_labels is not None and end_labels is not None and self.training:
402
+ # named entity detection loss
403
+
404
+ # start
405
+ if ned_start_logits is not None:
406
+ ned_start_loss = self.criterion(
407
+ ned_start_logits.view(-1, ned_start_logits.shape[-1]),
408
+ ned_start_labels.view(-1),
409
+ )
410
+ else:
411
+ ned_start_loss = 0
412
+
413
+ # end
414
+ # use ents_count to assign the labels to the correct positions i.e. using end_labels -> [[0,0,4,0], [0,0,0,2]] -> [4,2] (this is just an element, for batch we need to mask it with ents_count), ie -> [[4,2,-100,-100], [3,1,2,-100], [1,3,2,5]]
415
+
416
+ if ned_end_logits is not None:
417
+ ed_labels = end_labels.clone()
418
+ ed_labels = torch.nn.utils.rnn.pad_sequence(
419
+ torch.split(ed_labels[ed_labels > 0], ents_count.tolist()),
420
+ batch_first=True,
421
+ padding_value=-100,
422
+ )
423
+ end_labels[end_labels > 0] = 1
424
+ if not self.config.binary_end_logits:
425
+ # transform label to position in the sequence
426
+ end_labels = end_labels.argmax(dim=-1)
427
+ ned_end_loss = self.criterion(
428
+ ned_end_logits.view(-1, ned_end_logits.shape[-1]),
429
+ end_labels.view(-1),
430
+ )
431
+ else:
432
+ ned_end_loss = self.criterion(ned_end_logits.reshape(-1, ned_end_logits.shape[-1]), end_labels.reshape(-1).long())
433
+
434
+ # entity disambiguation loss
435
+ ed_loss = self.criterion(
436
+ ed_logits.view(-1, ed_logits.shape[-1]),
437
+ ed_labels.view(-1).long(),
438
+ )
439
+
440
+ else:
441
+ ned_end_loss = 0
442
+ ed_loss = 0
443
+
444
+ output_dict["ned_start_loss"] = ned_start_loss
445
+ output_dict["ned_end_loss"] = ned_end_loss
446
+ output_dict["ed_loss"] = ed_loss
447
+
448
+ output_dict["loss"] = ned_start_loss + ned_end_loss + ed_loss
449
+
450
+ return output_dict
451
+
452
+
453
+ class RelikReaderREModel(PreTrainedModel):
454
+ config_class = RelikReaderConfig
455
+
456
+ def __init__(self, config, *args, **kwargs):
457
+ super().__init__(config)
458
+ # Transformer model declaration
459
+ # self.transformer_model_name = transformer_model
460
+ self.config = config
461
+ self.transformer_model = (
462
+ AutoModel.from_pretrained(config.transformer_model)
463
+ if config.num_layers is None
464
+ else AutoModel.from_pretrained(
465
+ config.transformer_model, num_hidden_layers=config.num_layers
466
+ )
467
+ )
468
+ self.transformer_model.resize_token_embeddings(
469
+ self.transformer_model.config.vocab_size
470
+ + config.additional_special_symbols
471
+ + config.additional_special_symbols_types,
472
+ )
473
+
474
+ # named entity detection layers
475
+ self.ned_start_classifier = self._get_projection_layer(
476
+ config.activation, last_hidden=2, layer_norm=False
477
+ )
478
+
479
+ self.ned_end_classifier = PoolerEndLogitsBi(self.transformer_model.config)
480
+
481
+ self.relation_disambiguation_loss = (
482
+ config.relation_disambiguation_loss
483
+ if hasattr(config, "relation_disambiguation_loss")
484
+ else False
485
+ )
486
+
487
+ if self.config.entity_type_loss and self.config.add_entity_embedding:
488
+ input_hidden_ents = 3
489
+ else:
490
+ input_hidden_ents = 2
491
+
492
+ self.re_projector = self._get_projection_layer(
493
+ config.activation,
494
+ input_hidden=input_hidden_ents * self.transformer_model.config.hidden_size,
495
+ hidden=input_hidden_ents * self.config.linears_hidden_size,
496
+ last_hidden=2 * self.config.linears_hidden_size,
497
+ )
498
+
499
+ self.re_relation_projector = self._get_projection_layer(
500
+ config.activation,
501
+ input_hidden=self.transformer_model.config.hidden_size,
502
+ )
503
+
504
+ if self.config.entity_type_loss or self.relation_disambiguation_loss:
505
+ self.re_entities_projector = self._get_projection_layer(
506
+ config.activation,
507
+ input_hidden=2 * self.transformer_model.config.hidden_size,
508
+ )
509
+ self.re_definition_projector = self._get_projection_layer(
510
+ config.activation,
511
+ )
512
+
513
+ self.re_classifier = self._get_projection_layer(
514
+ config.activation,
515
+ input_hidden=config.linears_hidden_size,
516
+ last_hidden=2,
517
+ layer_norm=False,
518
+ )
519
+
520
+ self.training = config.training
521
+
522
+ # criterion
523
+ self.criterion = torch.nn.CrossEntropyLoss()
524
+ self.criterion_type = torch.nn.BCEWithLogitsLoss()
525
+
526
+ def _get_projection_layer(
527
+ self,
528
+ activation: str,
529
+ last_hidden: Optional[int] = None,
530
+ hidden: Optional[int] = None,
531
+ input_hidden=None,
532
+ layer_norm: bool = True,
533
+ ) -> torch.nn.Sequential:
534
+ head_components = [
535
+ torch.nn.Dropout(0.1),
536
+ torch.nn.Linear(
537
+ (
538
+ self.transformer_model.config.hidden_size
539
+ * self.config.use_last_k_layers
540
+ if input_hidden is None
541
+ else input_hidden
542
+ ),
543
+ self.config.linears_hidden_size if hidden is None else hidden,
544
+ ),
545
+ activation2functions[activation],
546
+ torch.nn.Dropout(0.1),
547
+ torch.nn.Linear(
548
+ self.config.linears_hidden_size if hidden is None else hidden,
549
+ self.config.linears_hidden_size if last_hidden is None else last_hidden,
550
+ ),
551
+ ]
552
+
553
+ if layer_norm:
554
+ head_components.append(
555
+ torch.nn.LayerNorm(
556
+ (
557
+ self.config.linears_hidden_size
558
+ if last_hidden is None
559
+ else last_hidden
560
+ ),
561
+ self.transformer_model.config.layer_norm_eps,
562
+ )
563
+ )
564
+
565
+ return torch.nn.Sequential(*head_components)
566
+
567
+ def _mask_logits(self, logits: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
568
+ mask = mask.unsqueeze(-1)
569
+ if next(self.parameters()).dtype == torch.float16:
570
+ logits = logits * (1 - mask) - 65500 * mask
571
+ else:
572
+ logits = logits * (1 - mask) - 1e30 * mask
573
+ return logits
574
+
575
+ def _get_model_features(
576
+ self,
577
+ input_ids: torch.Tensor,
578
+ attention_mask: torch.Tensor,
579
+ token_type_ids: Optional[torch.Tensor],
580
+ ):
581
+ model_input = {
582
+ "input_ids": input_ids,
583
+ "attention_mask": attention_mask,
584
+ "output_hidden_states": self.config.use_last_k_layers > 1,
585
+ }
586
+
587
+ if token_type_ids is not None:
588
+ model_input["token_type_ids"] = token_type_ids
589
+
590
+ model_output = self.transformer_model(**model_input)
591
+
592
+ if self.config.use_last_k_layers > 1:
593
+ model_features = torch.cat(
594
+ model_output[1][-self.config.use_last_k_layers :], dim=-1
595
+ )
596
+ else:
597
+ model_features = model_output[0]
598
+
599
+ return model_features
600
+
601
+ def compute_ned_end_logits(
602
+ self,
603
+ start_predictions,
604
+ start_labels,
605
+ model_features,
606
+ prediction_mask,
607
+ batch_size,
608
+ mask_preceding: bool = False,
609
+ ) -> Optional[torch.Tensor]:
610
+ # todo: maybe when constraining on the spans,
611
+ # we should not use a prediction_mask for the end tokens.
612
+ # at least we should not during training imo
613
+ start_positions = start_labels if self.training else start_predictions
614
+ start_positions_indices = (
615
+ torch.arange(start_positions.size(1), device=start_positions.device)
616
+ .unsqueeze(0)
617
+ .expand(batch_size, -1)[start_positions > 0]
618
+ ).to(start_positions.device)
619
+
620
+ if len(start_positions_indices) > 0:
621
+ expanded_features = model_features.repeat_interleave(
622
+ torch.sum(start_positions > 0, dim=-1), dim=0
623
+ )
624
+ expanded_prediction_mask = prediction_mask.repeat_interleave(
625
+ torch.sum(start_positions > 0, dim=-1), dim=0
626
+ )
627
+ if mask_preceding:
628
+ expanded_prediction_mask[
629
+ torch.arange(
630
+ expanded_prediction_mask.shape[1],
631
+ device=expanded_prediction_mask.device,
632
+ )
633
+ < start_positions_indices.unsqueeze(1)
634
+ ] = 1
635
+ end_logits = self.ned_end_classifier(
636
+ hidden_states=expanded_features,
637
+ start_positions=start_positions_indices,
638
+ p_mask=expanded_prediction_mask,
639
+ )
640
+
641
+ return end_logits
642
+
643
+ return None
644
+
645
+ def compute_relation_logits(
646
+ self,
647
+ model_entity_features,
648
+ special_symbols_features,
649
+ ) -> torch.Tensor:
650
+ model_subject_object_features = self.re_projector(model_entity_features)
651
+ model_subject_features = model_subject_object_features[
652
+ :, :, : model_subject_object_features.shape[-1] // 2
653
+ ]
654
+ model_object_features = model_subject_object_features[
655
+ :, :, model_subject_object_features.shape[-1] // 2 :
656
+ ]
657
+ special_symbols_start_representation = self.re_relation_projector(
658
+ special_symbols_features
659
+ )
660
+ re_logits = torch.einsum(
661
+ "bse,bde,bfe->bsdfe",
662
+ model_subject_features,
663
+ model_object_features,
664
+ special_symbols_start_representation,
665
+ )
666
+ re_logits = self.re_classifier(re_logits)
667
+
668
+ return re_logits
669
+
670
+ def compute_entity_logits(
671
+ self,
672
+ model_entity_features,
673
+ special_symbols_features,
674
+ ) -> torch.Tensor:
675
+ model_ed_features = self.re_entities_projector(model_entity_features)
676
+ special_symbols_ed_representation = self.re_definition_projector(
677
+ special_symbols_features
678
+ )
679
+
680
+ logits = torch.bmm(
681
+ model_ed_features,
682
+ torch.permute(special_symbols_ed_representation, (0, 2, 1)),
683
+ )
684
+ logits = self._mask_logits(
685
+ logits, (model_entity_features == -100).all(2).long()
686
+ )
687
+ return logits
688
+
689
+ def compute_loss(self, logits, labels, mask=None):
690
+ logits = logits.reshape(-1, logits.shape[-1])
691
+ labels = labels.reshape(-1).long()
692
+ if mask is not None:
693
+ return self.criterion(logits[mask], labels[mask])
694
+ return self.criterion(logits, labels)
695
+
696
+ def compute_ned_type_loss(
697
+ self,
698
+ disambiguation_labels,
699
+ re_ned_entities_logits,
700
+ ned_type_logits,
701
+ re_entities_logits,
702
+ entity_types,
703
+ mask,
704
+ ):
705
+ if self.config.entity_type_loss and self.relation_disambiguation_loss:
706
+ return self.criterion_type(
707
+ re_ned_entities_logits[disambiguation_labels != -100],
708
+ disambiguation_labels[disambiguation_labels != -100],
709
+ )
710
+ if self.config.entity_type_loss:
711
+ return self.criterion_type(
712
+ ned_type_logits[mask],
713
+ disambiguation_labels[:, :, :entity_types][mask],
714
+ )
715
+
716
+ if self.relation_disambiguation_loss:
717
+ return self.criterion_type(
718
+ re_entities_logits[disambiguation_labels != -100],
719
+ disambiguation_labels[disambiguation_labels != -100],
720
+ )
721
+ return 0
722
+
723
+ def compute_relation_loss(self, relation_labels, re_logits):
724
+ return self.compute_loss(
725
+ re_logits, relation_labels, relation_labels.view(-1) != -100
726
+ )
727
+
728
+ def forward(
729
+ self,
730
+ input_ids: torch.Tensor,
731
+ attention_mask: torch.Tensor,
732
+ token_type_ids: torch.Tensor,
733
+ prediction_mask: Optional[torch.Tensor] = None,
734
+ special_symbols_mask: Optional[torch.Tensor] = None,
735
+ special_symbols_mask_entities: Optional[torch.Tensor] = None,
736
+ start_labels: Optional[torch.Tensor] = None,
737
+ end_labels: Optional[torch.Tensor] = None,
738
+ disambiguation_labels: Optional[torch.Tensor] = None,
739
+ relation_labels: Optional[torch.Tensor] = None,
740
+ relation_threshold: float = None,
741
+ is_validation: bool = False,
742
+ is_prediction: bool = False,
743
+ use_predefined_spans: bool = False,
744
+ *args,
745
+ **kwargs,
746
+ ) -> Dict[str, Any]:
747
+ relation_threshold = (
748
+ self.config.threshold if relation_threshold is None else relation_threshold
749
+ )
750
+
751
+ batch_size = input_ids.shape[0]
752
+
753
+ model_features = self._get_model_features(
754
+ input_ids, attention_mask, token_type_ids
755
+ )
756
+
757
+ # named entity detection
758
+ if use_predefined_spans:
759
+ ned_start_logits, ned_start_probabilities, ned_start_predictions = (
760
+ None,
761
+ None,
762
+ torch.zeros_like(start_labels),
763
+ )
764
+ ned_end_logits, ned_end_probabilities, ned_end_predictions = (
765
+ None,
766
+ None,
767
+ torch.zeros_like(end_labels),
768
+ )
769
+
770
+ ned_start_predictions[start_labels > 0] = 1
771
+ ned_end_predictions[end_labels > 0] = 1
772
+ ned_end_predictions = ned_end_predictions[~(end_labels == -100).all(2)]
773
+ ned_start_labels = start_labels
774
+ ned_start_labels[start_labels > 0] = 1
775
+ else:
776
+ # start boundary prediction
777
+ ned_start_logits = self.ned_start_classifier(model_features)
778
+ if is_validation or is_prediction:
779
+ ned_start_logits = self._mask_logits(
780
+ ned_start_logits, prediction_mask
781
+ ) # why?
782
+ ned_start_probabilities = torch.softmax(ned_start_logits, dim=-1)
783
+ ned_start_predictions = ned_start_probabilities.argmax(dim=-1)
784
+
785
+ # end boundary prediction
786
+ ned_start_labels = (
787
+ torch.zeros_like(start_labels) if start_labels is not None else None
788
+ )
789
+
790
+ # start_labels contain entity id at their position, we just need 1 for start of entity
791
+ if ned_start_labels is not None:
792
+ ned_start_labels[start_labels == -100] = -100
793
+ ned_start_labels[start_labels > 0] = 1
794
+
795
+ # compute end logits only if there are any start predictions.
796
+ # For each start prediction, n end predictions are made
797
+ ned_end_logits = self.compute_ned_end_logits(
798
+ ned_start_predictions,
799
+ ned_start_labels,
800
+ model_features,
801
+ prediction_mask,
802
+ batch_size,
803
+ True,
804
+ )
805
+
806
+ if ned_end_logits is not None:
807
+ # For each start prediction, n end predictions are made based on
808
+ # binary classification ie. argmax at each position.
809
+ ned_end_probabilities = torch.softmax(ned_end_logits, dim=-1)
810
+ ned_end_predictions = ned_end_probabilities.argmax(dim=-1)
811
+ else:
812
+ ned_end_logits, ned_end_probabilities = None, None
813
+ ned_end_predictions = torch.zeros_like(ned_start_predictions)
814
+
815
+ if is_prediction or is_validation:
816
+ end_preds_count = ned_end_predictions.sum(1)
817
+ # If there are no end predictions for a start prediction, remove the start prediction
818
+ if (end_preds_count == 0).any() and (ned_start_predictions > 0).any():
819
+ ned_start_predictions[ned_start_predictions == 1] = (
820
+ end_preds_count != 0
821
+ ).long()
822
+ ned_end_predictions = ned_end_predictions[end_preds_count != 0]
823
+
824
+ if end_labels is not None:
825
+ end_labels = end_labels[~(end_labels == -100).all(2)]
826
+
827
+ start_position, end_position = (
828
+ (start_labels, end_labels)
829
+ if (not is_prediction and not is_validation)
830
+ else (ned_start_predictions, ned_end_predictions)
831
+ )
832
+
833
+ start_counts = (start_position > 0).sum(1)
834
+ if (start_counts > 0).any():
835
+ ned_end_predictions = ned_end_predictions.split(start_counts.tolist())
836
+ # limit to 30 predictions per document using start_counts, by setting all po after sum is 30 to 0
837
+ # if is_validation or is_prediction:
838
+ # ned_start_predictions[ned_start_predictions == 1] = start_counts
839
+ # We can only predict relations if we have start and end predictions
840
+ if (end_position > 0).sum() > 0:
841
+ ends_count = (end_position > 0).sum(1)
842
+ model_subject_features = torch.cat(
843
+ [
844
+ torch.repeat_interleave(
845
+ model_features[start_position > 0], ends_count, dim=0
846
+ ), # start position features
847
+ torch.repeat_interleave(model_features, start_counts, dim=0)[
848
+ end_position > 0
849
+ ], # end position features
850
+ ],
851
+ dim=-1,
852
+ )
853
+ ents_count = torch.nn.utils.rnn.pad_sequence(
854
+ torch.split(ends_count, start_counts.tolist()),
855
+ batch_first=True,
856
+ padding_value=0,
857
+ ).sum(1)
858
+ model_subject_features = torch.nn.utils.rnn.pad_sequence(
859
+ torch.split(model_subject_features, ents_count.tolist()),
860
+ batch_first=True,
861
+ padding_value=-100,
862
+ )
863
+
864
+ # if is_validation or is_prediction:
865
+ # model_subject_features = model_subject_features[:, :30, :]
866
+
867
+ # entity disambiguation. Here relation_disambiguation_loss would only be useful to
868
+ # reduce the number of candidate relations for the next step, but currently unused.
869
+ if self.config.entity_type_loss or self.relation_disambiguation_loss:
870
+ (re_ned_entities_logits) = self.compute_entity_logits(
871
+ model_subject_features,
872
+ model_features[
873
+ special_symbols_mask | special_symbols_mask_entities
874
+ ].view(batch_size, -1, model_features.shape[-1]),
875
+ )
876
+ entity_types = torch.sum(special_symbols_mask_entities, dim=1)[0].item()
877
+ ned_type_logits = re_ned_entities_logits[:, :, :entity_types]
878
+ re_entities_logits = re_ned_entities_logits[:, :, entity_types:]
879
+
880
+ if self.config.entity_type_loss:
881
+ ned_type_probabilities = torch.sigmoid(ned_type_logits)
882
+ ned_type_predictions = ned_type_probabilities.argmax(dim=-1)
883
+
884
+ if self.config.add_entity_embedding:
885
+ special_symbols_representation = model_features[
886
+ special_symbols_mask_entities
887
+ ].view(batch_size, entity_types, -1)
888
+
889
+ entities_representation = torch.einsum(
890
+ "bsp,bpe->bse",
891
+ ned_type_probabilities,
892
+ special_symbols_representation,
893
+ )
894
+ model_subject_features = torch.cat(
895
+ [model_subject_features, entities_representation], dim=-1
896
+ )
897
+ re_entities_probabilities = torch.sigmoid(re_entities_logits)
898
+ re_entities_predictions = re_entities_probabilities.round()
899
+ else:
900
+ (
901
+ ned_type_logits,
902
+ ned_type_probabilities,
903
+ re_entities_logits,
904
+ re_entities_probabilities,
905
+ ) = (None, None, None, None)
906
+ ned_type_predictions, re_entities_predictions = (
907
+ torch.zeros([batch_size, 1], dtype=torch.long).to(input_ids.device),
908
+ torch.zeros([batch_size, 1], dtype=torch.long).to(input_ids.device),
909
+ )
910
+
911
+ # Compute relation logits
912
+ re_logits = self.compute_relation_logits(
913
+ model_subject_features,
914
+ model_features[special_symbols_mask].view(
915
+ batch_size, -1, model_features.shape[-1]
916
+ ),
917
+ )
918
+
919
+ re_probabilities = torch.softmax(re_logits, dim=-1)
920
+ # we set a thresshold instead of argmax in cause it needs to be tweaked
921
+ re_predictions = re_probabilities[:, :, :, :, 1] > relation_threshold
922
+ re_probabilities = re_probabilities[:, :, :, :, 1]
923
+ else:
924
+ (
925
+ ned_type_logits,
926
+ ned_type_probabilities,
927
+ re_entities_logits,
928
+ re_entities_probabilities,
929
+ ) = (None, None, None, None)
930
+ ned_type_predictions, re_entities_predictions = (
931
+ torch.zeros([batch_size, 1], dtype=torch.long).to(input_ids.device),
932
+ torch.zeros([batch_size, 1], dtype=torch.long).to(input_ids.device),
933
+ )
934
+ re_logits, re_probabilities, re_predictions = (
935
+ torch.zeros(
936
+ [batch_size, 1, 1, special_symbols_mask.sum(1)[0]], dtype=torch.long
937
+ ).to(input_ids.device),
938
+ torch.zeros(
939
+ [batch_size, 1, 1, special_symbols_mask.sum(1)[0]], dtype=torch.long
940
+ ).to(input_ids.device),
941
+ torch.zeros(
942
+ [batch_size, 1, 1, special_symbols_mask.sum(1)[0]], dtype=torch.long
943
+ ).to(input_ids.device),
944
+ )
945
+
946
+ # output build
947
+ output_dict = dict(
948
+ batch_size=batch_size,
949
+ ned_start_logits=ned_start_logits,
950
+ ned_start_probabilities=ned_start_probabilities,
951
+ ned_start_predictions=ned_start_predictions,
952
+ ned_end_logits=ned_end_logits,
953
+ ned_end_probabilities=ned_end_probabilities,
954
+ ned_end_predictions=ned_end_predictions,
955
+ ned_type_logits=ned_type_logits,
956
+ ned_type_probabilities=ned_type_probabilities,
957
+ ned_type_predictions=ned_type_predictions,
958
+ re_entities_logits=re_entities_logits,
959
+ re_entities_probabilities=re_entities_probabilities,
960
+ re_entities_predictions=re_entities_predictions,
961
+ re_logits=re_logits,
962
+ re_probabilities=re_probabilities,
963
+ re_predictions=re_predictions,
964
+ )
965
+
966
+ if (
967
+ start_labels is not None
968
+ and end_labels is not None
969
+ and relation_labels is not None
970
+ and is_prediction is False
971
+ ):
972
+ ned_start_loss = self.compute_loss(ned_start_logits, ned_start_labels)
973
+ end_labels[end_labels > 0] = 1
974
+ ned_end_loss = self.compute_loss(ned_end_logits, end_labels)
975
+ if self.config.entity_type_loss or self.relation_disambiguation_loss:
976
+ ned_type_loss = self.compute_ned_type_loss(
977
+ disambiguation_labels,
978
+ re_ned_entities_logits,
979
+ ned_type_logits,
980
+ re_entities_logits,
981
+ entity_types,
982
+ (model_subject_features != -100).all(2),
983
+ )
984
+ relation_loss = self.compute_relation_loss(relation_labels, re_logits)
985
+ # compute loss. We can skip the relation loss if we are in the first epochs (optional)
986
+ if self.config.entity_type_loss or self.relation_disambiguation_loss:
987
+ output_dict["loss"] = (
988
+ ned_start_loss + ned_end_loss + relation_loss + ned_type_loss
989
+ ) / 4
990
+ output_dict["ned_type_loss"] = ned_type_loss
991
+ else:
992
+ output_dict["loss"] = ((1 / 20) * (ned_start_loss + ned_end_loss)) + (
993
+ (9 / 10) * relation_loss
994
+ )
995
+ output_dict["ned_start_loss"] = ned_start_loss
996
+ output_dict["ned_end_loss"] = ned_end_loss
997
+ output_dict["re_loss"] = relation_loss
998
+
999
+ return output_dict