Spaces:
Running on Zero
Running on Zero
File size: 5,849 Bytes
6b3a024 | 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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | # V-SPLADE
# Copyright (c) 2026-present NAVER Corp.
# Apache-2.0
"""
Query encoder strategies.
V_SPLADE quality variant uses an inference-free *Learned Sparse Retriever*
query encoder (``li_lsr``): a frozen copy of the backbone word embeddings,
optionally adapted with a low-rank LoRA, followed by a 1-dim projection
and an activation. At inference time the projection becomes a
``vocab_size``-entry lookup table, so encoding a query is just an
``Embedding(input_ids)`` — no transformer forward.
The simpler ``bow`` encoder turns the query into a binary token-presence
indicator over the vocab. Useful for smoke tests and the cheapest
inference-time setting.
"""
import torch
import torch.nn as nn
from enum import Enum
from typing import Optional
class QueryEncoderType(Enum):
BOW = "bow"
LI_LSR = "li_lsr"
class InferenceFreeQueryEncoder(nn.Module):
"""Li-LSR style inference-free sparse query encoder.
Training: token_id → embedding (frozen + LoRA) → projection(D→1) →
activation → scatter into a vocab-dim vector.
Inference: token_id → precomputed lookup table (no forward pass).
"""
def __init__(
self,
embed_weight: torch.Tensor,
vocab_size: int,
hidden_size: int,
lora_r: int = 0,
activation: str = "relu",
):
super().__init__()
assert activation in ("relu", "softplus"), f"Unknown activation: {activation}"
self.activation = activation
# Frozen copy of backbone word embeddings.
self.embeddings = nn.Embedding(vocab_size, hidden_size)
self.embeddings.weight.data.copy_(embed_weight.detach())
self.embeddings.weight.requires_grad = False
# Optional LoRA on the embedding lookup.
self.use_lora = lora_r > 0
if self.use_lora:
self.lora_A = nn.Parameter(torch.randn(vocab_size, lora_r) * 0.01)
self.lora_B = nn.Parameter(torch.zeros(lora_r, hidden_size))
# Projection: hidden_size → 1 (fully tuned).
self.projection = nn.Linear(hidden_size, 1, bias=True)
self.vocab_size = vocab_size
self._lookup_table: Optional[torch.Tensor] = None
def _activate(self, x: torch.Tensor) -> torch.Tensor:
if self.activation == "relu":
return torch.log1p(torch.relu(x))
return torch.nn.functional.softplus(x)
def _valid_mask(self, input_ids: torch.Tensor) -> torch.Tensor:
return input_ids < self.vocab_size
def _get_embed(self, input_ids: torch.Tensor) -> torch.Tensor:
safe_ids = input_ids.clamp(max=self.vocab_size - 1)
e = self.embeddings(safe_ids)
if self.use_lora:
e = e + self.lora_A[safe_ids] @ self.lora_B
return e
def forward(
self, input_ids: torch.Tensor, attention_mask: torch.Tensor,
) -> torch.Tensor:
"""Training forward. Returns (B, vocab_size) sparse vector."""
valid = self._valid_mask(input_ids)
mask = attention_mask.float() * valid.float()
scores = self.projection(self._get_embed(input_ids)).squeeze(-1)
scores = self._activate(scores) * mask
B = input_ids.size(0)
safe_ids = input_ids.clamp(max=self.vocab_size - 1)
out = torch.zeros(B, self.vocab_size, device=input_ids.device, dtype=scores.dtype)
out.scatter_add_(1, safe_ids, scores)
return out
@torch.no_grad()
def build_lookup_table(self) -> torch.Tensor:
device = self.embeddings.weight.device
all_ids = torch.arange(self.vocab_size, device=device)
e = self._get_embed(all_ids.unsqueeze(0))
s = self.projection(e).squeeze(-1).squeeze(0)
self._lookup_table = self._activate(s)
return self._lookup_table
def encode_with_lookup(
self, input_ids: torch.Tensor, attention_mask: torch.Tensor,
) -> torch.Tensor:
"""Inference: pure lookup, no embedding/projection forward."""
if self._lookup_table is None:
self.build_lookup_table()
valid = self._valid_mask(input_ids)
mask = attention_mask.float() * valid.float()
safe_ids = input_ids.clamp(max=self.vocab_size - 1)
scores = self._lookup_table.to(input_ids.device)[safe_ids] * mask
B = input_ids.size(0)
out = torch.zeros(B, self.vocab_size, device=input_ids.device, dtype=scores.dtype)
out.scatter_add_(1, safe_ids, scores)
return out
class BOWQueryEncoder(nn.Module):
"""Bag-of-Words: binary token indicator over the vocab."""
def __init__(self, vocab_size: int):
super().__init__()
self.vocab_size = vocab_size
def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
B, _ = input_ids.shape
device = input_ids.device
bow = torch.zeros(B, self.vocab_size, device=device, dtype=torch.bfloat16)
safe_ids = input_ids.clamp(0, self.vocab_size - 1) * attention_mask.long()
bow.scatter_(1, safe_ids, 1.0)
bow[:, 0] = 0.0 # clear padding token (id=0) accumulation
return bow
def build_query_encoder(
query_encoder_type: str,
vocab_size: int,
hidden_size: int = 768,
embed_weight: Optional[torch.Tensor] = None,
lora_r: int = 0,
activation: str = "relu",
**kwargs,
) -> nn.Module:
if query_encoder_type == "bow":
return BOWQueryEncoder(vocab_size)
if query_encoder_type == "li_lsr":
if embed_weight is None:
raise ValueError("li_lsr requires embed_weight from the encoder")
return InferenceFreeQueryEncoder(
embed_weight, vocab_size, hidden_size,
lora_r=lora_r, activation=activation,
)
raise ValueError(
f"Unknown query_encoder_type: {query_encoder_type}. Choose: bow | li_lsr"
)
|