| """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 |
| 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]) |
|
|
| 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)]) |
| return tok |
|
|
