|
|
| import json, numpy as np, torch |
| import torch.nn as nn |
| from safetensors.torch import load_file |
|
|
| class TransformerBlock(nn.Module): |
| def __init__(self, dim, heads, ff_dim, dropout): |
| super().__init__() |
| self.mha = nn.MultiheadAttention(dim, heads, dropout=dropout, batch_first=True) |
| self.ln1 = nn.LayerNorm(dim); self.ln2 = nn.LayerNorm(dim) |
| self.drop = nn.Dropout(dropout) |
| self.ffn = nn.Sequential(nn.Linear(dim, ff_dim), nn.ReLU(), nn.Dropout(dropout), nn.Linear(ff_dim, dim)) |
| def forward(self, x, key_padding_mask=None): |
| attn_out, _ = self.mha(x, x, x, key_padding_mask=key_padding_mask, need_weights=False) |
| x = self.ln1(x + self.drop(attn_out)) |
| ff_out = self.ffn(x) |
| return self.ln2(x + self.drop(ff_out)) |
|
|
| class BiGRUResidualBlock(nn.Module): |
| def __init__(self, dim, dropout): |
| super().__init__() |
| self.gru = nn.GRU(dim, dim//2, num_layers=1, batch_first=True, bidirectional=True) |
| self.ln = nn.LayerNorm(dim); self.drop = nn.Dropout(dropout) |
| def forward(self, x): |
| out, _ = self.gru(x) |
| return self.ln(x + self.drop(out)) |
|
|
| class MeterModel(nn.Module): |
| def __init__(self, vocab_size, num_classes, T, emb_dim=64, num_heads=4, ff_dim=256, gru_blocks=3, dropout=0.1): |
| super().__init__() |
| self.emb = nn.Embedding(vocab_size, emb_dim, padding_idx=0) |
| self.pos = nn.Embedding(T, emb_dim) |
| self.drop = nn.Dropout(dropout) |
| self.tr = TransformerBlock(emb_dim, num_heads, ff_dim, dropout) |
| self.gru_blocks = nn.ModuleList([BiGRUResidualBlock(emb_dim, dropout) for _ in range(gru_blocks)]) |
| self.head = nn.Sequential(nn.Linear(emb_dim, 128), nn.ReLU(), nn.Dropout(dropout), nn.Linear(128, num_classes)) |
| def forward(self, x): |
| B, T = x.shape |
| pos = torch.arange(T, device=x.device).unsqueeze(0).expand(B, T) |
| h = self.drop(self.emb(x) + self.pos(pos)) |
| pad_mask = (x == 0) |
| h = self.tr(h, key_padding_mask=pad_mask) |
| for blk in self.gru_blocks: |
| h = blk(h) |
| mask = (~pad_mask).float().unsqueeze(-1) |
| pooled = (h * mask).sum(1) / mask.sum(1).clamp_min(1.0) |
| return self.head(pooled) |
|
|
| def load_local(repo_dir="."): |
| cfg = json.load(open(f"{repo_dir}/config.json", "r", encoding="utf-8")) |
| stoi = json.load(open(f"{repo_dir}/stoi.json", "r", encoding="utf-8")) |
| classes = cfg["classes"] |
| model = MeterModel( |
| vocab_size=cfg["vocab_size"], num_classes=cfg["num_classes"], T=cfg["max_length"], |
| emb_dim=cfg["emb_dim"], num_heads=cfg["num_heads"], ff_dim=cfg["ff_dim"], |
| gru_blocks=cfg["gru_blocks"], dropout=cfg["dropout"] |
| ) |
| sd = load_file(f"{repo_dir}/model.safetensors") |
| model.load_state_dict(sd) |
| model.eval() |
| return model, stoi, classes, cfg["max_length"] |
|
|
| def encode(text, stoi, max_len): |
| ids = np.zeros((max_len,), dtype=np.int64) |
| for i, ch in enumerate(text[:max_len]): |
| ids[i] = stoi.get(ch, 0) |
| return torch.from_numpy(ids).unsqueeze(0) |
|
|
| @torch.no_grad() |
| def predict(text, model, stoi, classes, max_len, topk=5): |
| x = encode(text, stoi, max_len) |
| logits = model(x)[0] |
| probs = torch.softmax(logits, dim=-1) |
| topv, topi = torch.topk(probs, k=min(topk, probs.numel())) |
| pred = classes[int(topi[0])] |
| top = [(classes[int(i)], float(v)) for v, i in zip(topv, topi)] |
| return pred, top |
|
|
