Views
No views yet
answerdotai/ModernBERT-large. The model is designed to perform two distinct text classification tasks using a shared feature representation, enhanced by a Mixture-of-Experts (MoE) layer.transformers library. The following code demonstrates how to make a prediction:1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4from transformers import AutoTokenizer, AutoModel
5from huggingface_hub import PyTorchModelHubMixin
6
7class MoELayer(nn.Module):
8 def __init__(self, input_dim, num_experts, top_k=2):
9 super(MoELayer, self).__init__()
10 self.num_experts = num_experts
11 self.top_k = top_k
12
13 # Define experts as independent feed-forward layers
14 self.experts = nn.ModuleList([nn.Sequential(
15 nn.Linear(input_dim, input_dim * 2),
16 nn.ReLU(),
17 nn.Linear(input_dim * 2, input_dim)
18 ) for _ in range(num_experts)])
19
20 self.gating_network = nn.Linear(input_dim, num_experts)
21
22 def forward(self, x):
23 gate_logits = self.gating_network(x)
24 gate_probs = F.softmax(gate_logits, dim=-1)
25
26 # Get top-k experts for each input
27 topk_vals, topk_indices = torch.topk(gate_probs, self.top_k, dim=-1)
28
29 # Compute contributions from top-k experts
30 output = torch.zeros_like(x)
31 for i in range(self.top_k):
32 expert_idx = topk_indices[:, i]
33 expert_weight = topk_vals[:, i].unsqueeze(-1)
34
35 expert_outputs = torch.stack([self.experts[j](x[b]) for b, j in enumerate(expert_idx)], dim=0)
36
37 output += expert_weight * expert_outputs
38
39 return output
40
41class SentenceClassificationMoeMTLModel(
42 nn.Module,
43 PyTorchModelHubMixin,
44):
45 def __init__(self) -> None:
46 super(SentenceClassificationMoeMTLModel, self).__init__()
47 self.base_model = AutoModel.from_pretrained("answerdotai/ModernBERT-large")
48
49 self.moe_layer = MoELayer(input_dim=self.base_model.config.hidden_size, num_experts=8, top_k=2)
50
51 self.task_1_classifier = nn.Sequential(
52 nn.Linear(in_features=self.base_model.config.hidden_size, out_features=768, bias=False),
53 nn.GELU(),
54 nn.LayerNorm(768, eps=1e-05, elementwise_affine=True),
55 nn.Linear(768, 2)
56 )
57
58 self.task_2_classifier = nn.Sequential(
59 nn.Linear(in_features=self.base_model.config.hidden_size, out_features=768, bias=False),
60 nn.GELU(),
61 nn.LayerNorm(768, eps=1e-05, elementwise_affine=True),
62 nn.Linear(768, 2),
63 )
64
65 def forward(self, task, input_ids, attention_mask):
66 x = self.base_model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
67 cls_r = x[:, 0]
68
69 x = self.moe_layer(x[:, 0])
70
71 if task == "arg":
72 x = self.task_1_classifier(x)
73 elif task == "stance":
74 x = self.task_2_classifier(x)
75
76 return x, cls_r
77
78model_name = "ag-charalampous/argument-same-side-stance-classification"
79tokenizer = AutoTokenizer.from_pretrained(model_name)
80
81model = SentenceClassificationMoeMTLModel.from_pretrained(model_name)
82model.eval()
83
84device = "cpu"
85
86def classify_sequence(seq, task, label_map):
87 enc = tokenizer(
88 *(seq if task == 'stance' else (seq,)),
89 return_tensors="pt",
90 truncation=True,
91 max_length=1024
92 ).to(device)
93
94 with torch.no_grad():
95 logits, _ = model(task=task, **enc)
96 probs = torch.softmax(logits, dim=-1).squeeze()
97 pred_idx = probs.argmax().item()
98 confidence = probs[pred_idx].item()
99
100 return label_map[pred_idx], confidence
101
102# Example input for task 1
103text = "A fetus or embryo is not a person; therefore, abortion should not be considered murder."
104
105label_map = {0: "Non-argumentative", 1: "Argumentative"}
106label, confidence = classify_sequence(text, 'arg', label_map)
107
108print(f"Prediction: {label} (Confidence: {confidence:.2f})")
109
110# Example input for task 2
111claim_1 = "A fetus or embryo is not a person; therefore, abortion should not be considered murder."
112claim_2 = "Since death is the intention, such procedures should be considered murder."
113
114label_map = {0: "Same-side", 1: "Opposing-side"}
115label, confidence = classify_sequence([claim_1, claim_2], 'stance', label_map)
116
117print(f"Prediction: {label} (Confidence: {confidence:.2f})")