| import torch |
| import torch.nn as nn |
| from transformers import AutoModel |
|
|
| class NeoAraBERTEntityPairTwoHead(nn.Module): |
| |
|
|
| def __init__( |
| self, |
| model_name, |
| num_positive_relations, |
| no_relation_id, |
| original_id_to_positive_id, |
| type_dim=42, |
| pair_hidden_1=1024, |
| pair_hidden_2=512, |
| dropout_p=0.1, |
| existence_loss_weight=1.0, |
| positive_relation_loss_weight=1.0, |
| ): |
| super().__init__() |
| self.encoder = AutoModel.from_pretrained(model_name, trust_remote_code=True) |
| hidden_size = self.encoder.config.hidden_size |
| self.no_relation_id = int(no_relation_id) |
| self.existence_loss_weight = float(existence_loss_weight) |
| self.positive_relation_loss_weight = float(positive_relation_loss_weight) |
| self.register_buffer( |
| "original_id_to_positive_id", |
| torch.tensor(original_id_to_positive_id, dtype=torch.long), |
| ) |
| self.pair_feature_dim = 5 * hidden_size + type_dim |
| self.pair_mlp = nn.Sequential( |
| nn.Linear(self.pair_feature_dim, pair_hidden_1), |
| nn.LayerNorm(pair_hidden_1), |
| nn.GELU(), |
| nn.Dropout(dropout_p), |
| nn.Linear(pair_hidden_1, pair_hidden_2), |
| nn.LayerNorm(pair_hidden_2), |
| nn.GELU(), |
| nn.Dropout(dropout_p), |
| ) |
| self.existence_head = nn.Linear(pair_hidden_2, 2) |
| self.positive_relation_head = nn.Linear(pair_hidden_2, num_positive_relations) |
|
|
| def forward( |
| self, |
| input_ids, |
| attention_mask, |
| type_features, |
| subject_marker_positions, |
| object_marker_positions, |
| labels=None, |
| ): |
| hidden = self.encoder( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| ).last_hidden_state |
| batch_indices = torch.arange(hidden.shape[0], device=hidden.device) |
| cls = hidden[:, 0, :] |
| subject = hidden[batch_indices, subject_marker_positions.long(), :] |
| object_ = hidden[batch_indices, object_marker_positions.long(), :] |
| pair_features = torch.cat([ |
| cls, |
| subject, |
| object_, |
| torch.abs(subject - object_), |
| subject * object_, |
| type_features.float(), |
| ], dim=-1) |
| pair_hidden = self.pair_mlp(pair_features) |
| existence_logits = self.existence_head(pair_hidden) |
| positive_logits = self.positive_relation_head(pair_hidden) |
| packed_logits = torch.cat([existence_logits, positive_logits], dim=-1) |
| loss = None |
| if labels is not None: |
| existence_targets = (labels != self.no_relation_id).long() |
| existence_loss = nn.functional.cross_entropy(existence_logits, existence_targets) |
| positive_mask = existence_targets == 1 |
| if positive_mask.any(): |
| positive_targets = self.original_id_to_positive_id[labels[positive_mask]] |
| positive_loss = nn.functional.cross_entropy( |
| positive_logits[positive_mask], |
| positive_targets, |
| ) |
| else: |
| positive_loss = existence_loss.new_zeros(()) |
| loss = ( |
| self.existence_loss_weight * existence_loss |
| + self.positive_relation_loss_weight * positive_loss |
| ) |
| result = {"logits": packed_logits} |
| if loss is not None: |
| result["loss"] = loss |
| return result |
|
|
