| |
| """Self-contained Kenga instruct encoder. Byte UTF-8, bidirectional. |
| |
| Usage (retrieval, Giga-style prefix on the query only): |
| |
| from modeling_kenga_embed import KengaEmbed |
| m = KengaEmbed.from_pretrained("GermannM/kenga-embed-prophet-instruct") |
| q = m.encode_queries(["Where is the capital of Russia?"]) |
| p = m.encode_documents(["Moscow is the capital of Russia."]) |
| |
| STS (prefix on both sentences): |
| |
| a = m.encode_sts(["The cat sits on the mat."]) |
| b = m.encode_sts(["A cat is sitting on a mat."]) |
| """ |
| import json |
| import os |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| INSTR_RETR = "Instruct: Given a query, retrieve relevant passages\nQuery: " |
| INSTR_STS = "Instruct: Retrieve semantically similar text.\nQuery: " |
|
|
|
|
| class BL(nn.Module): |
| def __init__(self, in_f, out_f, k): |
| super().__init__() |
| kk = min(k, in_f, out_f) |
| self.U = nn.Parameter(torch.zeros(out_f, kk)) |
| self.V = nn.Parameter(torch.zeros(in_f, kk)) |
| self.S = nn.Parameter(torch.ones(kk)) |
| self.b = nn.Parameter(torch.zeros(out_f)) |
| self.register_buffer("mask", torch.ones(kk)) |
|
|
| def set_active(self, k_t): |
| m = torch.zeros_like(self.mask) |
| m[: min(int(k_t), self.mask.numel())] = 1 |
| self.mask.copy_(m) |
|
|
| def forward(self, x): |
| return ((x @ self.V) * (self.S * self.mask)) @ self.U.T + self.b |
|
|
|
|
| class BAttn(nn.Module): |
| def __init__(self, d, heads, k_attn): |
| super().__init__() |
| self.h = heads |
| self.dh = d // heads |
| self.scale = self.dh ** -0.5 |
| self.q = BL(d, d, k_attn) |
| self.k_ = BL(d, d, k_attn) |
| self.v = BL(d, d, k_attn) |
| self.o = BL(d, d, k_attn) |
|
|
| def set_rank(self, ka): |
| for bl in (self.q, self.k_, self.v, self.o): |
| bl.set_active(ka) |
|
|
| def forward(self, x, pad): |
| B, T, _ = x.shape |
| Q = self.q(x).view(B, T, self.h, self.dh).transpose(1, 2) |
| K = self.k_(x).view(B, T, self.h, self.dh).transpose(1, 2) |
| V = self.v(x).view(B, T, self.h, self.dh).transpose(1, 2) |
| att = (Q.float() @ K.float().transpose(-2, -1)) * self.scale |
| att = att.masked_fill(pad, float("-inf")).softmax(-1) |
| return self.o((att @ V).transpose(1, 2).reshape(B, T, x.size(-1))) |
|
|
|
|
| class BBlock(nn.Module): |
| def __init__(self, d, dff, heads, k_attn, k_ff): |
| super().__init__() |
| self.ln1 = nn.LayerNorm(d) |
| self.ln2 = nn.LayerNorm(d) |
| self.attn = BAttn(d, heads, k_attn) |
| self.fc1 = BL(d, dff, k_ff) |
| self.fc2 = BL(dff, d, k_ff) |
|
|
| def set_rank(self, ka, kf): |
| self.attn.set_rank(ka) |
| self.fc1.set_active(kf) |
| self.fc2.set_active(kf) |
|
|
| def forward(self, x, pad): |
| x = x + self.attn(self.ln1(x), pad) |
| return x + self.fc2(F.gelu(self.fc1(self.ln2(x)))) |
|
|
|
|
| class KengaEmbedEnc(nn.Module): |
| def __init__(self, d=768, layers=8, heads=12, dff=3072, seq=160, |
| v=256, k_attn=192, k_ff=512): |
| super().__init__() |
| self.d = d |
| self.seq = seq |
| self.tok = nn.Embedding(v, d) |
| self.pos = nn.Parameter(torch.zeros(1, seq, d)) |
| self.blocks = nn.ModuleList( |
| [BBlock(d, dff, heads, k_attn, k_ff) for _ in range(layers)] |
| ) |
| self.ln = nn.LayerNorm(d) |
| self.proj = nn.Linear(d, d) |
|
|
| def set_rank(self, ka, kf): |
| for b in self.blocks: |
| b.set_rank(ka, kf) |
|
|
| def forward(self, idx, lengths=None): |
| B, T = idx.shape |
| x = self.tok(idx) + self.pos[:, :T] |
| pad = idx.eq(0) |
| mask = pad[:, None, None, :].expand(B, 1, T, T) |
| for blk in self.blocks: |
| x = blk(x, mask) |
| x = self.ln(x) |
| m = (~pad).unsqueeze(-1).float() |
| pooled = (x * m).sum(1) / m.sum(1).clamp(min=1) |
| return F.normalize(self.proj(pooled), dim=-1) |
|
|
|
|
| class KengaEmbed: |
| def __init__(self, model, seq, instruct=True, device="cpu"): |
| self.model = model.to(device).eval() |
| self.seq = seq |
| self.instruct = instruct |
| self.device = device |
|
|
| @classmethod |
| def from_pretrained(cls, path, device="cpu"): |
| cfg = json.load(open(os.path.join(path, "config.json"), encoding="utf-8")) |
| m = KengaEmbedEnc( |
| d=cfg.get("d", 768), |
| layers=cfg.get("layers", 8), |
| heads=cfg.get("heads", 12), |
| dff=cfg.get("dff", 3072), |
| seq=cfg.get("seq", 160), |
| v=cfg.get("v", 256), |
| k_attn=cfg.get("k_attn", 192), |
| k_ff=cfg.get("k_ff", 512), |
| ) |
| blob = torch.load(os.path.join(path, "pytorch_model.bin"), |
| map_location="cpu", weights_only=False) |
| sd = blob["model"] if isinstance(blob, dict) and "model" in blob else blob |
| m.load_state_dict(sd, strict=True) |
| rank = cfg.get("rank") or blob.get("rank") |
| if rank: |
| m.set_rank(*list(rank)) |
| return cls(m, seq=cfg.get("seq", 160), |
| instruct=bool(cfg.get("instruct", True)), device=device) |
|
|
| def _ids(self, text): |
| raw = (text or "").encode("utf-8", errors="ignore")[: self.seq] |
| ids = [b if b else 1 for b in raw][: self.seq] |
| n = len(ids) |
| ids = ids + [0] * (self.seq - n) |
| return ids, n |
|
|
| def _prefix(self, text, kind): |
| if not self.instruct or not text: |
| return text |
| if text.startswith("Instruct:"): |
| return text |
| p = INSTR_STS if kind == "sts" else INSTR_RETR |
| return p + text |
|
|
| def encode(self, texts, kind="doc", bs=32): |
| vecs = [] |
| with torch.no_grad(): |
| for i in range(0, len(texts), bs): |
| chunk = texts[i : i + bs] |
| ids, ln = [], [] |
| for t in chunk: |
| s = t if kind == "doc" else self._prefix(t, kind) |
| ii, n = self._ids(s) |
| ids.append(ii) |
| ln.append(n) |
| x = torch.tensor(ids, device=self.device) |
| l = torch.tensor(ln, device=self.device) |
| vecs.append(self.model(x, l).cpu()) |
| return torch.cat(vecs, 0) |
|
|
| def encode_queries(self, texts, **kw): |
| return self.encode(texts, kind="retr", **kw) |
|
|
| def encode_documents(self, texts, **kw): |
| return self.encode(texts, kind="doc", **kw) |
|
|
| def encode_sts(self, texts, **kw): |
| return self.encode(texts, kind="sts", **kw) |
|
|
