| """ChoiceEncoder: choice label -> semantic embedding in the core space. |
| |
| Choices are *inputs*: "yes", "no", "database", "move up" are encoded from their |
| text, so the same weights serve 2, 4 or 10 options and unseen labels. With a |
| pretrained text encoder the choice embedding comes from the same backbone as the |
| question (shared weights), so "sì" and "yes" start close together. |
| |
| `EmbeddingCache` memoises backbone outputs per string. With a frozen backbone a |
| question or a choice label is encoded once, ever: Snake's 4 moves and a fixed |
| question cost nothing at inference, only the core runs. |
| """ |
| from __future__ import annotations |
|
|
| from collections import OrderedDict |
|
|
| import torch |
| from torch import nn |
|
|
| from minijev.encoders.text import Projection, TextPooler |
|
|
|
|
| class ChoiceEncoder(nn.Module): |
| def __init__(self, d: int, embed: nn.Embedding | None = None, max_len: int = 16, d_text: int | None = None): |
| super().__init__() |
| self.byte = embed is not None |
| if self.byte: |
| self.pooler = TextPooler(embed, d, max_len) |
| else: |
| self.proj = Projection(d_text, d, hidden=2 * d) |
|
|
| def forward(self, ids=None, pooled=None): |
| """byte: ids (..., L) -> (..., d); pretrained: backbone pooled (..., d_text) -> (..., d).""" |
| return self.pooler(ids) if self.byte else self.proj(pooled) |
|
|
|
|
| class EmbeddingCache: |
| """LRU cache: text -> backbone output (tokens, mask, pooled) as CPU tensors.""" |
|
|
| def __init__(self, max_items: int = 4096): |
| self.max_items = max_items |
| self.data: OrderedDict[str, tuple] = OrderedDict() |
| self.hits = self.misses = 0 |
|
|
| def get_many(self, texts, compute): |
| """compute(list[str]) -> (tok (n, L, dt), mask (n, L), pooled (n, dt)) for the missing texts.""" |
| missing = [t for t in dict.fromkeys(texts) if t not in self.data] |
| self.hits += len(texts) - len(missing) |
| self.misses += len(missing) |
| if missing: |
| tok, mask, pooled = compute(missing) |
| for i, t in enumerate(missing): |
| n = int(mask[i].sum()) |
| self.data[t] = (tok[i, :n].detach().cpu(), pooled[i].detach().cpu()) |
| if len(self.data) > self.max_items: |
| self.data.popitem(last=False) |
| for t in texts: |
| self.data.move_to_end(t) |
| rows = [self.data[t] for t in texts] |
| L = max(1, max(r[0].shape[0] for r in rows)) |
| dt = rows[0][1].shape[-1] |
| tok = torch.zeros(len(rows), L, dt) |
| mask = torch.zeros(len(rows), L, dtype=torch.bool) |
| for i, (t, _) in enumerate(rows): |
| tok[i, : t.shape[0]] = t |
| mask[i, : t.shape[0]] = True |
| return tok, mask, torch.stack([r[1] for r in rows]) |
|
|
