File size: 7,055 Bytes
d4337b2 | 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 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | # coding=utf-8
"""
A sentence-transformers input module that reproduces the LLM2Vec encoding
pipeline exactly: instruction handling, left-padded tokenization, and a
mean pooling that covers only the content tokens (BOS and instruction
tokens are excluded from the pooled vector).
This is a faithful port of `llm2vec.LLM2Vec.encode` -> the model produced
here yields the same embeddings as the original llm2vec-based usage, but
loads with a plain `SentenceTransformer(...)` call and needs no extra
package beyond `sentence-transformers` + `transformers`.
"""
from __future__ import annotations
from typing import Any
import torch
from transformers import AutoModel, AutoTokenizer
from sentence_transformers.base.modules.input_module import InputModule
class LLM2VecTransformer(InputModule):
"""LLM2Vec encoder packaged as a sentence-transformers module."""
# Internal separator used by llm2vec to split instruction from content.
SEP = "!@#$%^&*()"
config_file_name: str = "sentence_bert_config.json"
config_keys: list[str] = ["max_seq_length", "pooling_mode", "model_kwargs"]
save_in_root: bool = True
def __init__(
self,
model_name_or_path: str,
max_seq_length: int = 8124,
pooling_mode: str = "mean",
model_kwargs: dict[str, Any] | None = None,
tokenizer_kwargs: dict[str, Any] | None = None,
hub_kwargs: dict[str, Any] | None = None,
**kwargs,
) -> None:
super().__init__()
if pooling_mode != "mean":
raise ValueError("Only 'mean' pooling is supported by this model.")
self.max_seq_length = max_seq_length
self.pooling_mode = pooling_mode
# `model_kwargs` is persisted in the module config; `hub_kwargs`
# (token, revision, trust_remote_code, ...) are loading-only.
self.model_kwargs = model_kwargs or {}
hub_kwargs = hub_kwargs or {}
self.auto_model = AutoModel.from_pretrained(
model_name_or_path, **{**self.model_kwargs, **hub_kwargs}
)
self.tokenizer = AutoTokenizer.from_pretrained(
model_name_or_path, **{**(tokenizer_kwargs or {}), **hub_kwargs}
)
# llm2vec configures the tokenizer this way before encoding.
self.tokenizer.pad_token = self.tokenizer.eos_token
self.tokenizer.padding_side = "left"
# -- encoding -----------------------------------------------------------
def preprocess(
self,
inputs: list[str],
prompt: str | None = None,
**kwargs,
) -> dict[str, torch.Tensor | Any]:
"""Tokenize, reproducing llm2vec `_convert_to_str` + `tokenize`."""
sep = self.SEP
# llm2vec format: "<instruction> !@#$%^&*()<text>" for queries,
# "!@#$%^&*()<text>" for documents (no instruction).
if prompt:
combined = [f"{prompt.strip()} {sep}{text}" for text in inputs]
else:
combined = [f"{sep}{text}" for text in inputs]
# Split off the content (everything after the separator).
content = [t.split(sep)[1] if len(t.split(sep)) > 1 else "" for t in combined]
originals = ["".join(t.split(sep)) for t in combined]
features = self.tokenizer(
originals,
return_tensors="pt",
padding=True,
truncation=True,
max_length=self.max_seq_length,
)
# `embed_mask` marks the content tokens (the trailing tokens of each
# sequence), excluding BOS, instruction and padding.
embed_mask = torch.zeros_like(features["attention_mask"])
for i, text in enumerate(content):
ids = self.tokenizer(
[text],
return_tensors="pt",
padding=True,
truncation=True,
max_length=self.max_seq_length,
add_special_tokens=False,
)["input_ids"][0]
if len(ids) > 0:
embed_mask[i, -len(ids):] = 1
features["embed_mask"] = embed_mask
return features
def forward(self, features: dict[str, Any], **kwargs) -> dict[str, Any]:
"""Run the encoder and mean-pool over content tokens only."""
outputs = self.auto_model(
input_ids=features["input_ids"],
attention_mask=features["attention_mask"],
)
last_hidden = outputs.last_hidden_state
embed_mask = features["embed_mask"]
# Identical to llm2vec's mean pooling (left-padded, content at the end).
seq_lengths = embed_mask.sum(dim=-1)
embeddings = torch.stack(
[
last_hidden[i, -int(length):, :].mean(dim=0)
for i, length in enumerate(seq_lengths)
],
dim=0,
)
features["sentence_embedding"] = embeddings
return features
# -- introspection ------------------------------------------------------
def get_sentence_embedding_dimension(self) -> int:
return self.auto_model.config.hidden_size
@property
def max_seq_length_(self) -> int: # pragma: no cover - convenience
return self.max_seq_length
# -- persistence --------------------------------------------------------
def save(self, output_path: str, *args, safe_serialization: bool = True, **kwargs) -> None:
self.auto_model.save_pretrained(output_path, safe_serialization=safe_serialization)
self.tokenizer.save_pretrained(output_path)
self.save_config(output_path)
@classmethod
def load(
cls,
model_name_or_path: str,
subfolder: str = "",
token: bool | str | None = None,
cache_folder: str | None = None,
revision: str | None = None,
local_files_only: bool = False,
trust_remote_code: bool = False,
model_kwargs: dict[str, Any] | None = None,
processor_kwargs: dict[str, Any] | None = None,
config_kwargs: dict[str, Any] | None = None,
backend: str = "torch",
**kwargs,
) -> "LLM2VecTransformer":
config = cls.load_config(
model_name_or_path=model_name_or_path,
subfolder=subfolder,
token=token,
cache_folder=cache_folder,
revision=revision,
local_files_only=local_files_only,
)
hub_kwargs = {
"token": token,
"cache_dir": cache_folder,
"revision": revision,
"local_files_only": local_files_only,
"trust_remote_code": trust_remote_code,
}
persistent_model_kwargs = {
**config.get("model_kwargs", {}),
**(model_kwargs or {}),
}
return cls(
model_name_or_path,
max_seq_length=config.get("max_seq_length", 8124),
pooling_mode=config.get("pooling_mode", "mean"),
model_kwargs=persistent_model_kwargs,
tokenizer_kwargs=processor_kwargs or {},
hub_kwargs=hub_kwargs,
)
|