File size: 3,343 Bytes
7ed86c3 | 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 | """MeterModel: Stoicheia encoder + per-letter macron and scansion heads."""
from __future__ import annotations
from dataclasses import dataclass, field
import torch
import torch.nn as nn
import torch.nn.functional as F
from meter.backbone import CharBertWithHidden
@dataclass
class MeterConfig:
head_dropout: float = 0.33
w_mac: float = 1.0 # 0 disables the task (head still built, frozen out of loss)
w_scan: float = 1.0
use_cap: bool = True # zero-init capitalization input embedding
scalar_mix: bool = True # ELMo-style learned mix over all block outputs
mac_class_w: list = field(default_factory=lambda: [1.0, 1.0]) # long, short
scan_class_w: list = field(default_factory=lambda: [1.0, 1.0, 1.0, 1.0])
class MeterModel(nn.Module):
def __init__(self, encoder: CharBertWithHidden, mcfg: MeterConfig):
super().__init__()
self.encoder = encoder
self.mcfg = mcfg
d = encoder.cfg.d_model
if mcfg.use_cap:
emb = nn.Embedding(2, d)
nn.init.zeros_(emb.weight)
encoder.cap_emb = emb # picked up by CharBertWithHidden.forward
self.dropout = nn.Dropout(mcfg.head_dropout)
self.head_mac = nn.Linear(d, 2, bias=False)
self.head_scan = nn.Linear(d, 4, bias=False)
if mcfg.scalar_mix:
encoder.return_layers = True
# blocks + final normed hidden; zero-init = uniform mix at start
self.mix_w = nn.Parameter(torch.zeros(len(encoder.blocks) + 1))
for m in (self.head_mac, self.head_scan):
nn.init.normal_(m.weight, std=0.02)
self.register_buffer("mac_w", torch.tensor(mcfg.mac_class_w, dtype=torch.float32))
self.register_buffer("scan_w", torch.tensor(mcfg.scan_class_w, dtype=torch.float32))
# pretraining output heads take no part in the loss; freeze them so DDP
# doesn't trip on parameters that never receive gradients
for m in (encoder.head_char, encoder.head_bnd, encoder.head_dia,
encoder.head_cap, encoder.head_punct):
for p in m.parameters():
p.requires_grad_(False)
def forward(self, batch):
out = self.encoder(batch)
if self.mcfg.scalar_mix:
h = torch.stack([*out["layers"], out["hidden"]]) # (L+1,B,T,D)
mix = torch.softmax(self.mix_w, 0)
h = torch.einsum("l,lbtd->btd", mix.to(h.dtype), h)
else:
h = out["hidden"]
h = self.dropout(h)
return dict(mac=self.head_mac(h), scan=self.head_scan(h))
def _ce(self, logits, target, weight):
"""Class-weighted CE that stays finite (and keeps the head in the DDP graph)
when a batch has no valid labels for this task."""
if bool((target != -100).any()):
return F.cross_entropy(logits.transpose(1, 2), target,
weight=weight.to(logits.dtype), ignore_index=-100)
return logits.sum() * 0.0
def loss(self, out, batch):
t = self.mcfg
l_m = self._ce(out["mac"], batch["y_mac"], self.mac_w)
l_s = self._ce(out["scan"], batch["y_scan"], self.scan_w)
loss = t.w_mac * l_m + t.w_scan * l_s
return loss, dict(l_mac=round(l_m.item(), 4), l_scan=round(l_s.item(), 4))
|