vector-index-embedding 0.1.0__py3-none-any.whl

vector_index_embedding-0.1.0.dist-info/METADATA

@@ -0,0 +1,28 @@
+Metadata-Version: 2.4
+Name: vector-index-embedding
+Version: 0.1.0
+Summary: HNSW based output embeddings for LLM's
+Author-email: Martin Loretz <pypi@martinloretz.com>
+License-Expression: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: torch
+Requires-Dist: hnswlib
+Requires-Dist: huggingface_hub
+
+# vector-index-embedding
+
+
+
+
+## Faster hnswlib
+
+Our faster implementation can be found here: https://github.com/martinloretzzz/hnswlib
+
+Warning: This implementation might not work on all systems as it was only tested on the one where we'Re running the bnechmarks and the SIMD implementation was only adapted for that architecture.
+
+For benchmarking we use our own fork of hnswlib, that includes 2 improvements for fast inner product distances on high dimensional data:
+- We calculate all the inner products in paralell, that way we reduce memory accesses in half (we load one element of the query and compare it to N other vectors at the same time)
+- We removed a heuristic that restricted multi-threading, as as our data is extremly high dimesional and always benifit from using all cores.

vector_index_embedding-0.1.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+vectorindex/__init__.py,sha256=QcAEBigbm7RZxX8u4F6hrB85tuS77IMCrbBgJ-6dOtE,118
+vectorindex/flat_index.py,sha256=ouTpOBhuT28pOpQXLbyRzMYdAPu2SXy-h0IOfIATe4g,588
+vectorindex/vector_index.py,sha256=e9ViG9N6sEbbLXkbS-BHxqrEy7d_pzcgzMksTnpa67c,4725
+vector_index_embedding-0.1.0.dist-info/METADATA,sha256=sUGd7MZAoS3878UvHZMdTmP5j1QTWuA9tQbBCBwJICc,1215
+vector_index_embedding-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+vector_index_embedding-0.1.0.dist-info/top_level.txt,sha256=3UTzrAMuIGn-k6iQ_K6X2x51t8eNAp6VKHD3-VyOCbM,12
+vector_index_embedding-0.1.0.dist-info/RECORD,,

vector_index_embedding-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

vector_index_embedding-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+vectorindex

vectorindex/__init__.py

@@ -0,0 +1,2 @@
+from .vector_index import VectorIndexEmbedding, VectorIndexEmbeddingConfig
+from .flat_index import FlatIndexEmbedding

vectorindex/flat_index.py

@@ -0,0 +1,16 @@
+import torch
+import torch.nn as nn
+
+class FlatIndexEmbedding(nn.Module):
+    def __init__(self, layer: nn.Linear, k: int):
+        super().__init__()
+        self.layer, self.k = layer, k
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.layer(x)
+        return self.mask_non_topk(x, k=self.k, fill=float("-inf"))
+
+    def mask_non_topk(self, x: torch.Tensor, k: int, fill: float):
+        _, indices = torch.topk(x, k, dim=-1)
+        mask = torch.full_like(x, fill, dtype=x.dtype, device=x.device)
+        return mask.scatter(-1, indices, x.gather(-1, indices))

vectorindex/vector_index.py

@@ -0,0 +1,102 @@
+from huggingface_hub import hf_hub_download
+from huggingface_hub.errors import RemoteEntryNotFoundError
+import torch
+import torch.nn as nn
+import hnswlib
+from pathlib import Path
+from dataclasses import dataclass, asdict
+import json
+
+@dataclass
+class VectorIndexEmbeddingConfig:
+    model_name: str
+    k: int
+    ef: int
+    M: int
+    ef_construction: int
+    special_tokens: list[int] | None = None
+    dim: int = -1 # set from weight in build_index
+    vocab_size: int = -1 # set from weight in build_index
+    model_id: str | None = None
+
+class VectorIndexEmbedding(nn.Module):
+    def __init__(self, config: VectorIndexEmbeddingConfig, index_path: str):
+        super().__init__()
+        self.config = config
+        self.index = hnswlib.Index(space='ip', dim=int(config.dim))
+        self.index.load_index(index_path)
+        self.index.set_ef(config.ef)
+        self.num_threads = -1
+
+        if config.special_tokens is not None:
+            self.special_token_indices = torch.tensor(config.special_tokens, dtype=torch.long)
+            self.special_token_weight = torch.from_numpy(self.index.get_items(config.special_tokens, return_type="numpy"))
+
+    @torch.compiler.disable
+    def topk(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        indices, distances = self.index.knn_query(x.detach().float().numpy(), k=self.config.k, num_threads=self.num_threads)
+        return 1.0 - torch.from_numpy(distances), torch.from_numpy(indices).to(torch.int64)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # return torch.full((x.shape[0], x.shape[1], self.config.vocab_size), 0, dtype=x.dtype, device=x.device)
+
+        x_flat = x.view(-1, x.shape[-1]).float()
+        distances, indices = self.topk(x_flat.cpu())
+
+        logits = torch.full((x_flat.shape[0], self.config.vocab_size), float("-inf"), dtype=x.dtype, device=x.device)
+        logits.scatter_(-1, indices.to(x.device), distances.to(x.device).to(x.dtype))
+
+        if self.config.special_tokens is not None:
+            special_token_distances = torch.matmul(x_flat, self.special_token_weight.to(x.device).T).to(x.dtype)
+            special_token_indices = self.special_token_indices.to(x.device).unsqueeze(0).expand(x_flat.shape[0], -1)
+            logits.scatter_(-1, special_token_indices, special_token_distances)
+
+        return logits.view((x.shape[0], x.shape[1], self.config.vocab_size))
+
+    def set_ef(self, ef: int):
+        self.index.set_ef(ef)
+
+    @staticmethod
+    def from_pretrained(model_id: str, ef = None, k = None, repo_id = "martinloretzzz/vector-index-embedding") -> "VectorIndexEmbedding":
+        try:
+            index_name = VectorIndexEmbedding.get_index_name(model_id)
+            local_path = hf_hub_download(repo_id=repo_id, filename=f"{index_name}.index")
+            local_path_config = hf_hub_download(repo_id=repo_id, filename=f"{index_name}.json")
+            return VectorIndexEmbedding.from_file(local_path, ef=ef, k=k, config_path=local_path_config)
+        except RemoteEntryNotFoundError:
+            raise Exception(f"No prebuilt vector index for model '{model_id}' was found. To build your own index, please follow the guide at: https://github.com/martinloretzzz/vector-index-embedding")
+
+    @staticmethod
+    def from_file(path: str, ef = None, k = None, config_path: str | None = None) -> "VectorIndexEmbedding":
+        index_path = Path(path)
+        if config_path is None:
+            config_path = index_path.resolve().with_suffix(".json")
+        with open(config_path, "r") as f:
+            config = VectorIndexEmbeddingConfig(**json.load(f))
+
+        if k is not None: config.k = k
+        if ef is not None: config.ef = ef 
+
+        return VectorIndexEmbedding(config, index_path=str(index_path.absolute()))
+
+    @staticmethod
+    def build_index(weight: torch.Tensor, config: VectorIndexEmbeddingConfig, save_path: str = "data", seed=42):
+        config.vocab_size, config.dim = weight.shape
+        index_file = Path(save_path) / Path(f"{config.model_name}.index") 
+        index_file.parent.mkdir(parents=True, exist_ok=True)
+
+        index = hnswlib.Index(space='ip', dim=config.dim)
+        index.init_index(max_elements=config.vocab_size, M=config.M, ef_construction=config.ef_construction, random_seed=seed)
+        index.set_ef(config.ef)
+        index.add_items(weight.cpu().numpy())
+        index.save_index(str(index_file))
+
+        with open(index_file.resolve().with_suffix(".json"), "w") as f:
+            json.dump(asdict(config), f, indent=4)
+
+        print(f"Index saved to {index_file}")
+        return str(index_file)
+
+    @staticmethod
+    def get_index_name(hf_model_id: str):
+        return hf_model_id.lower().replace("/", "-").replace(".", "-")