File size: 3,613 Bytes
ac20d60 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | import torch
import torch.nn as nn
from transformers import AutoModel
class NeoAraBERTEntityPairTwoHead(nn.Module):
# Entity-pair relation classifier with binary existence and 40-way positive heads.
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
|