File size: 2,747 Bytes
026f63e | 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 | """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])
|