Spaces:
Sleeping
Sleeping
File size: 2,976 Bytes
9ee1e28 |
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 |
import torch
import torch.nn as nn
class StanceClassifier(nn.Module):
def __init__(self, transformer_model, num_classes, dropout_rate=0.6):
super(StanceClassifier, self).__init__()
self.transformer = transformer_model
self.dropout = nn.Dropout(dropout_rate)
self.layer_norm = nn.LayerNorm(transformer_model.config.hidden_size)
l0 = transformer_model.config.hidden_size
l1 = transformer_model.config.hidden_size * 2
l2 = l1 // 2
l3 = l2 // 2
# classifier expects pooled token representation (batch, hidden)
self.classifier = nn.Sequential(
nn.Linear(l0, l1),
nn.LayerNorm(l1),
nn.GELU(),
nn.Dropout(dropout_rate),
nn.Linear(l1, l2),
nn.LayerNorm(l2),
nn.GELU(),
nn.Dropout(dropout_rate),
nn.Linear(l2, l3),
nn.LayerNorm(l3),
nn.GELU(),
nn.Linear(l3, num_classes),
)
self.attention_vector = nn.Linear(l0, 1)
nn.init.xavier_uniform_(self.attention_vector.weight)
self.freeze_transformer()
def freeze_transformer(self):
for param in self.transformer.parameters():
param.requires_grad = False
def unfreeze_transformer(self):
for param in self.transformer.parameters():
param.requires_grad = True
def forward(self, input_ids, attention_mask):
if not any(p.requires_grad for p in self.transformer.parameters()):
with torch.no_grad():
outputs = self.transformer(
input_ids=input_ids, attention_mask=attention_mask
)
else:
outputs = self.transformer(
input_ids=input_ids, attention_mask=attention_mask
)
# token-level hidden states: (batch, seq_len, hidden)
token_states = outputs.last_hidden_state
scores = self.attention_vector(token_states).squeeze(-1) # (batch, seq_len)
mask = attention_mask.to(dtype=torch.bool) # (batch, seq_len)
scores = scores.masked_fill(~mask, -1e9)
weights = torch.softmax(scores, dim=1) # (batch, seq_len)
pooled_output = (weights.unsqueeze(-1) * token_states).sum(dim=1) # (batch, hidden)
if torch.isnan(pooled_output).any() or torch.isinf(pooled_output).any():
print("WARNING: Transformer output NaN/Inf")
pooled_output = torch.where(
torch.isnan(pooled_output) | torch.isinf(pooled_output),
torch.zeros_like(pooled_output),
pooled_output,
)
pooled_output = self.layer_norm(pooled_output)
logits = self.classifier(pooled_output)
return logits
def classifier_params(self):
return list(self.classifier.parameters())
def transformer_params(self):
return list(self.transformer.parameters()) |