File size: 3,175 Bytes
e1679fa | 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 | """DeBERTa-v3 with a small head. One class, three jobs, chosen by `task`:
- "scores": six numbers per essay (mean pooling, linear head)
"""
import torch
import torch.nn as nn
from transformers import AutoConfig, AutoModel, AutoTokenizer
from coach.data import BR, TRAITS
BACKBONE = "microsoft/deberta-v3-base"
def tokenizer(name=BACKBONE):
tok = AutoTokenizer.from_pretrained(name)
tok.add_special_tokens({"additional_special_tokens": [BR]})
return tok
class ScoreModel(nn.Module):
def __init__(self, name=BACKBONE, n_tokens=None, n_out=len(TRAITS)):
super().__init__()
cfg = AutoConfig.from_pretrained(name)
cfg.hidden_dropout_prob = 0.0 # the ELL winners all switched dropout off for regression
cfg.attention_probs_dropout_prob = 0.0
self.backbone = AutoModel.from_pretrained(name, config=cfg)
if n_tokens:
self.backbone.resize_token_embeddings(n_tokens)
self.head = nn.Linear(cfg.hidden_size, n_out)
def forward(self, input_ids, attention_mask):
h = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
m = attention_mask.unsqueeze(-1).to(h.dtype)
pooled = (h * m).sum(1) / m.sum(1).clamp(min=1)
return self.head(pooled)
def param_groups(model, lr, head_lr, decay=0.9, weight_decay=0.01):
"""Layer-wise learning-rate decay: the head and top layer at `lr`, each layer below at lr * decay."""
groups = [{"params": list(model.head.parameters()), "lr": head_lr, "weight_decay": 0.0}]
layers = list(model.backbone.encoder.layer)
n = len(layers)
for i, layer in enumerate(layers):
groups.append({"params": list(layer.parameters()), "lr": lr * decay ** (n - 1 - i), "weight_decay": weight_decay})
emb = [p for n_, p in model.backbone.named_parameters() if not n_.startswith("encoder.layer")]
groups.append({"params": emb, "lr": lr * decay ** n, "weight_decay": weight_decay})
return groups
class PartsModel(nn.Module):
"""Token classification over the whole essay: which argument part each token belongs to (BIO over 7 types), and
how effective that part is (3 classes, read off the tokens of the part and averaged at decode time)."""
def __init__(self, name=BACKBONE, n_tokens=None, n_bio=15, n_eff=3):
super().__init__()
cfg = AutoConfig.from_pretrained(name)
self.backbone = AutoModel.from_pretrained(name, config=cfg)
if n_tokens:
self.backbone.resize_token_embeddings(n_tokens)
self.bio = nn.Linear(cfg.hidden_size, n_bio)
self.eff = nn.Linear(cfg.hidden_size, n_eff)
self.head = nn.ModuleList([self.bio, self.eff]) # so param_groups finds "head"
def forward(self, input_ids, attention_mask):
h = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
return self.bio(h), self.eff(h)
def parts_tokenizer(name=BACKBONE):
from tokenizers import AddedToken
tok = AutoTokenizer.from_pretrained(name)
tok.add_tokens([AddedToken("\n", normalized=False)]) # keep paragraph breaks; the tokenizer drops them otherwise
return tok
|