vortexa 0.3.1__tar.gz → 0.3.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (38) hide show
  1. {vortexa-0.3.1 → vortexa-0.3.2}/PKG-INFO +3 -7
  2. {vortexa-0.3.1 → vortexa-0.3.2}/README.md +2 -5
  3. {vortexa-0.3.1 → vortexa-0.3.2}/pyproject.toml +1 -2
  4. vortexa-0.3.2/src/vortexa/core/embedding.py +88 -0
  5. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/indexer.py +1 -1
  6. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa.egg-info/PKG-INFO +3 -7
  7. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa.egg-info/SOURCES.txt +0 -1
  8. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa.egg-info/requires.txt +0 -1
  9. vortexa-0.3.1/src/vortexa/core/embedding.py +0 -180
  10. vortexa-0.3.1/src/vortexa/core/lf4_model.py +0 -168
  11. {vortexa-0.3.1 → vortexa-0.3.2}/LICENSE +0 -0
  12. {vortexa-0.3.1 → vortexa-0.3.2}/setup.cfg +0 -0
  13. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/__init__.py +0 -0
  14. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/__init__.py +0 -0
  15. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/chunking.py +0 -0
  16. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/context_engine.py +0 -0
  17. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/graph.py +0 -0
  18. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/inference.py +0 -0
  19. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/language.py +0 -0
  20. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/lf4_v4_model.py +0 -0
  21. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/types.py +0 -0
  22. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/v4_embedder.py +0 -0
  23. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/vortex_score.py +0 -0
  24. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/interfaces/__init__.py +0 -0
  25. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/interfaces/cli.py +0 -0
  26. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/interfaces/mcp_server.py +0 -0
  27. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/interfaces/watcher.py +0 -0
  28. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/search/__init__.py +0 -0
  29. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/search/ranking.py +0 -0
  30. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/search/search.py +0 -0
  31. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/search/tokens.py +0 -0
  32. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/storage/__init__.py +0 -0
  33. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/storage/bm25.py +0 -0
  34. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/storage/vector_store.py +0 -0
  35. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/storage/walker.py +0 -0
  36. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa.egg-info/dependency_links.txt +0 -0
  37. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa.egg-info/entry_points.txt +0 -0
  38. {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: vortexa
3
- Version: 0.3.1
3
+ Version: 0.3.2
4
4
  Summary: Codebase indexing and semantic search engine
5
5
  Author-email: VortexAI <koulabhay25@gmail.com>
6
6
  License-Expression: Apache-2.0
@@ -29,7 +29,6 @@ Requires-Dist: huggingface-hub>=0.20.0
29
29
  Requires-Dist: tokenizers>=0.19.0
30
30
  Requires-Dist: safetensors>=0.4.0
31
31
  Requires-Dist: fastmcp>=2.0.0
32
- Requires-Dist: model2vec>=0.3.0
33
32
  Requires-Dist: sentence-transformers>=2.2.0
34
33
  Requires-Dist: tree-sitter-language-pack>=0.1.0
35
34
  Dynamic: license-file
@@ -375,11 +374,8 @@ indexer = CodebaseIndexer(root="/path/to/project", model_id="VTXAI/vtx-embed-7M"
375
374
 
376
375
  ### Alternative Embedders
377
376
 
378
- The `embedding.py` module also provides alternative embedders that can be passed directly:
379
-
380
- - `Model2VecEmbedder` — static embeddings via Model2Vec (`AI4free/JARVIS-tool-search-v1`)
381
- - `SentenceTransformerEmbedder` — Transformer-based dense embeddings (`all-MiniLM-L6-v2`)
382
- - `LF4Embedder` — 4-bit quantized static embeddings (`VTXAI/Vortex-Embed-4.7M`)
377
+ The `embedding.py` module also provides `SentenceTransformerEmbedder`
378
+ for transformer-based models (requires the `sentence-transformers` package).
383
379
 
384
380
  ---
385
381
 
@@ -339,11 +339,8 @@ indexer = CodebaseIndexer(root="/path/to/project", model_id="VTXAI/vtx-embed-7M"
339
339
 
340
340
  ### Alternative Embedders
341
341
 
342
- The `embedding.py` module also provides alternative embedders that can be passed directly:
343
-
344
- - `Model2VecEmbedder` — static embeddings via Model2Vec (`AI4free/JARVIS-tool-search-v1`)
345
- - `SentenceTransformerEmbedder` — Transformer-based dense embeddings (`all-MiniLM-L6-v2`)
346
- - `LF4Embedder` — 4-bit quantized static embeddings (`VTXAI/Vortex-Embed-4.7M`)
342
+ The `embedding.py` module also provides `SentenceTransformerEmbedder`
343
+ for transformer-based models (requires the `sentence-transformers` package).
347
344
 
348
345
  ---
349
346
 
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "vortexa"
7
- version = "0.3.1"
7
+ version = "0.3.2"
8
8
  description = "Codebase indexing and semantic search engine"
9
9
  readme = "README.md"
10
10
  requires-python = ">=3.10"
@@ -35,7 +35,6 @@ dependencies = [
35
35
  "tokenizers>=0.19.0",
36
36
  "safetensors>=0.4.0",
37
37
  "fastmcp>=2.0.0",
38
- "model2vec>=0.3.0",
39
38
  "sentence-transformers>=2.2.0",
40
39
  "tree-sitter-language-pack>=0.1.0",
41
40
  ]
@@ -0,0 +1,88 @@
1
+ """Embedding model abstraction for the codebase indexer.
2
+
3
+ Provides lazy-loading, thread-safe embedders with memoization support.
4
+ Inspired by cocoindex's SentenceTransformerEmbedder pattern.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ import logging
10
+ import threading
11
+ from typing import Protocol, runtime_checkable
12
+
13
+ import numpy as np
14
+ import numpy.typing as npt
15
+
16
+ logger = logging.getLogger(__name__)
17
+
18
+
19
+ @runtime_checkable
20
+ class Embedder(Protocol):
21
+ """Protocol for embedding models used by the indexer."""
22
+
23
+ @property
24
+ def dim(self) -> int:
25
+ """Embedding dimensionality."""
26
+ ...
27
+
28
+ def embed(self, text: str) -> npt.NDArray[np.float32]:
29
+ """Embed a single text string."""
30
+ ...
31
+
32
+ def embed_batch(self, texts: list[str]) -> npt.NDArray[np.float32]:
33
+ """Embed a batch of text strings."""
34
+ ...
35
+
36
+ @property
37
+ def memo_key(self) -> tuple:
38
+ """Identity key for memoization cache invalidation."""
39
+ ...
40
+
41
+
42
+ class SentenceTransformerEmbedder:
43
+ """Thread-safe embedder wrapping sentence-transformers.
44
+
45
+ Supports any sentence-transformers model with lazy loading.
46
+ Memo key includes model name and device for cache invalidation.
47
+ """
48
+
49
+ def __init__(self, model_name: str = "all-MiniLM-L6-v2", device: str | None = None) -> None:
50
+ self._model_name = model_name
51
+ self._device = device
52
+ self._model = None
53
+ self._lock = threading.Lock()
54
+
55
+ @property
56
+ def dim(self) -> int:
57
+ self._ensure_loaded()
58
+ assert self._model is not None
59
+ dim = self._model.get_embedding_dimension()
60
+ assert dim is not None
61
+ return dim
62
+
63
+ def _ensure_loaded(self) -> None:
64
+ if self._model is None:
65
+ with self._lock:
66
+ if self._model is None:
67
+ from sentence_transformers import SentenceTransformer
68
+ logger.info("Loading sentence-transformers model: %s", self._model_name)
69
+ self._model = SentenceTransformer(self._model_name, device=self._device)
70
+
71
+ def embed(self, text: str) -> npt.NDArray[np.float32]:
72
+ """Embed a single text string."""
73
+ self._ensure_loaded()
74
+ assert self._model is not None
75
+ return self._model.encode(text, convert_to_numpy=True, normalize_embeddings=True)
76
+
77
+ def embed_batch(self, texts: list[str]) -> npt.NDArray[np.float32]:
78
+ """Embed a batch of text strings."""
79
+ if not texts:
80
+ return np.empty((0, 0), dtype=np.float32)
81
+ self._ensure_loaded()
82
+ assert self._model is not None
83
+ return self._model.encode(texts, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=False)
84
+
85
+ @property
86
+ def memo_key(self) -> tuple:
87
+ """Identity key: (class, model_name, device)."""
88
+ return ("SentenceTransformerEmbedder", self._model_name, self._device)
@@ -21,7 +21,7 @@ from typing import cast
21
21
  import lmdb
22
22
  import numpy as np
23
23
  from vortexa.core.chunking import chunk_source
24
- from vortexa.core.embedding import Embedder, LF4Embedder
24
+ from vortexa.core.embedding import Embedder
25
25
  from vortexa.core.language import detect_language, get_extensions
26
26
  from vortexa.core.types import (
27
27
  Chunk,
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: vortexa
3
- Version: 0.3.1
3
+ Version: 0.3.2
4
4
  Summary: Codebase indexing and semantic search engine
5
5
  Author-email: VortexAI <koulabhay25@gmail.com>
6
6
  License-Expression: Apache-2.0
@@ -29,7 +29,6 @@ Requires-Dist: huggingface-hub>=0.20.0
29
29
  Requires-Dist: tokenizers>=0.19.0
30
30
  Requires-Dist: safetensors>=0.4.0
31
31
  Requires-Dist: fastmcp>=2.0.0
32
- Requires-Dist: model2vec>=0.3.0
33
32
  Requires-Dist: sentence-transformers>=2.2.0
34
33
  Requires-Dist: tree-sitter-language-pack>=0.1.0
35
34
  Dynamic: license-file
@@ -375,11 +374,8 @@ indexer = CodebaseIndexer(root="/path/to/project", model_id="VTXAI/vtx-embed-7M"
375
374
 
376
375
  ### Alternative Embedders
377
376
 
378
- The `embedding.py` module also provides alternative embedders that can be passed directly:
379
-
380
- - `Model2VecEmbedder` — static embeddings via Model2Vec (`AI4free/JARVIS-tool-search-v1`)
381
- - `SentenceTransformerEmbedder` — Transformer-based dense embeddings (`all-MiniLM-L6-v2`)
382
- - `LF4Embedder` — 4-bit quantized static embeddings (`VTXAI/Vortex-Embed-4.7M`)
377
+ The `embedding.py` module also provides `SentenceTransformerEmbedder`
378
+ for transformer-based models (requires the `sentence-transformers` package).
383
379
 
384
380
  ---
385
381
 
@@ -16,7 +16,6 @@ src/vortexa/core/graph.py
16
16
  src/vortexa/core/indexer.py
17
17
  src/vortexa/core/inference.py
18
18
  src/vortexa/core/language.py
19
- src/vortexa/core/lf4_model.py
20
19
  src/vortexa/core/lf4_v4_model.py
21
20
  src/vortexa/core/types.py
22
21
  src/vortexa/core/v4_embedder.py
@@ -6,6 +6,5 @@ huggingface-hub>=0.20.0
6
6
  tokenizers>=0.19.0
7
7
  safetensors>=0.4.0
8
8
  fastmcp>=2.0.0
9
- model2vec>=0.3.0
10
9
  sentence-transformers>=2.2.0
11
10
  tree-sitter-language-pack>=0.1.0
@@ -1,180 +0,0 @@
1
- """Embedding model abstraction for the codebase indexer.
2
-
3
- Provides lazy-loading, thread-safe embedders with memoization support.
4
- Inspired by cocoindex's SentenceTransformerEmbedder pattern.
5
- """
6
-
7
- from __future__ import annotations
8
-
9
- import logging
10
- import threading
11
- from typing import Protocol, runtime_checkable
12
-
13
- import numpy as np
14
- import numpy.typing as npt
15
-
16
- logger = logging.getLogger(__name__)
17
-
18
-
19
- @runtime_checkable
20
- class Embedder(Protocol):
21
- """Protocol for embedding models used by the indexer."""
22
-
23
- @property
24
- def dim(self) -> int:
25
- """Embedding dimensionality."""
26
- ...
27
-
28
- def embed(self, text: str) -> npt.NDArray[np.float32]:
29
- """Embed a single text string."""
30
- ...
31
-
32
- def embed_batch(self, texts: list[str]) -> npt.NDArray[np.float32]:
33
- """Embed a batch of text strings."""
34
- ...
35
-
36
- @property
37
- def memo_key(self) -> tuple:
38
- """Identity key for memoization cache invalidation."""
39
- ...
40
-
41
-
42
- class Model2VecEmbedder:
43
- """Thread-safe, lazy-loading embedder wrapping model2vec.StaticModel.
44
-
45
- The model is loaded on first use and cached. Thread-safe via a lock.
46
- Memo key includes the model ID for cache invalidation.
47
- """
48
-
49
- def __init__(self, model_id: str = "AI4free/JARVIS-tool-search-v1") -> None:
50
- self._model_id = model_id
51
- self._model = None
52
- self._lock = threading.Lock()
53
-
54
- @property
55
- def dim(self) -> int:
56
- self._ensure_loaded()
57
- assert self._model is not None
58
- return self._model.dim
59
-
60
- def _ensure_loaded(self) -> None:
61
- if self._model is None:
62
- with self._lock:
63
- if self._model is None: # Double-checked locking
64
- from model2vec import StaticModel
65
- logger.info("Loading embedding model: %s", self._model_id)
66
- self._model = StaticModel.from_pretrained(self._model_id)
67
-
68
- def embed(self, text: str) -> npt.NDArray[np.float32]:
69
- """Embed a single text string."""
70
- self._ensure_loaded()
71
- assert self._model is not None
72
- return self._model.encode([text])[0]
73
-
74
- def embed_batch(self, texts: list[str]) -> npt.NDArray[np.float32]:
75
- """Embed a batch of text strings."""
76
- if not texts:
77
- return np.empty((0, 0), dtype=np.float32)
78
- self._ensure_loaded()
79
- assert self._model is not None
80
- result = self._model.encode(texts)
81
- return np.array(result, dtype=np.float32)
82
-
83
- @property
84
- def memo_key(self) -> tuple:
85
- """Identity key: (class, model_id)."""
86
- return ("Model2VecEmbedder", self._model_id)
87
-
88
-
89
- class SentenceTransformerEmbedder:
90
- """Thread-safe embedder wrapping sentence-transformers.
91
-
92
- Supports any sentence-transformers model with lazy loading.
93
- Memo key includes model name and device for cache invalidation.
94
- """
95
-
96
- def __init__(self, model_name: str = "all-MiniLM-L6-v2", device: str | None = None) -> None:
97
- self._model_name = model_name
98
- self._device = device
99
- self._model = None
100
- self._lock = threading.Lock()
101
-
102
- @property
103
- def dim(self) -> int:
104
- self._ensure_loaded()
105
- assert self._model is not None
106
- dim = self._model.get_embedding_dimension()
107
- assert dim is not None
108
- return dim
109
-
110
- def _ensure_loaded(self) -> None:
111
- if self._model is None:
112
- with self._lock:
113
- if self._model is None:
114
- from sentence_transformers import SentenceTransformer
115
- logger.info("Loading sentence-transformers model: %s", self._model_name)
116
- self._model = SentenceTransformer(self._model_name, device=self._device)
117
-
118
- def embed(self, text: str) -> npt.NDArray[np.float32]:
119
- """Embed a single text string."""
120
- self._ensure_loaded()
121
- assert self._model is not None
122
- return self._model.encode(text, convert_to_numpy=True, normalize_embeddings=True)
123
-
124
- def embed_batch(self, texts: list[str]) -> npt.NDArray[np.float32]:
125
- """Embed a batch of text strings."""
126
- if not texts:
127
- return np.empty((0, 0), dtype=np.float32)
128
- self._ensure_loaded()
129
- assert self._model is not None
130
- return self._model.encode(texts, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=False)
131
-
132
- @property
133
- def memo_key(self) -> tuple:
134
- """Identity key: (class, model_name, device)."""
135
- return ("SentenceTransformerEmbedder", self._model_name, self._device)
136
-
137
-
138
- class LF4Embedder:
139
- """Thread-safe, lazy-loading embedder wrapping LF4StaticEmbedding (4-bit quantized).
140
-
141
- Uses the VTXAI/Vortex-Embed-4.7M model by default — a 4-bit static embedding
142
- model with ~3.5 MB footprint. Loads on first use, cached thereafter.
143
- """
144
-
145
- def __init__(self, model_id: str = "VTXAI/Vortex-Embed-4.7M") -> None:
146
- self._model_id = model_id
147
- self._model = None
148
- self._lock = threading.Lock()
149
-
150
- @property
151
- def dim(self) -> int:
152
- self._ensure_loaded()
153
- assert self._model is not None
154
- return self._model.dim
155
-
156
- def _ensure_loaded(self) -> None:
157
- if self._model is None:
158
- with self._lock:
159
- if self._model is None:
160
- logger.info("Loading LF4 embedding model: %s", self._model_id)
161
- from vortexa.core.lf4_model import LF4StaticEmbedding
162
- self._model = LF4StaticEmbedding.from_pretrained(self._model_id)
163
-
164
- def embed(self, text: str) -> npt.NDArray[np.float32]:
165
- """Embed a single text string."""
166
- self._ensure_loaded()
167
- assert self._model is not None
168
- return self._model.encode([text])[0]
169
-
170
- def embed_batch(self, texts: list[str]) -> npt.NDArray[np.float32]:
171
- """Embed a batch of text strings."""
172
- if not texts:
173
- return np.empty((0, 0), dtype=np.float32)
174
- self._ensure_loaded()
175
- assert self._model is not None
176
- return self._model.encode(texts)
177
-
178
- @property
179
- def memo_key(self) -> tuple:
180
- return ("LF4Embedder", self._model_id)
@@ -1,168 +0,0 @@
1
- """LF4 Static Embedding Model - Native 4-bit quantized sentence embeddings.
2
-
3
- Usage:
4
- from lf4_model import LF4StaticEmbedding
5
- model = LF4StaticEmbedding.from_pretrained("VTXAI/Vortex-Embed-4.7M")
6
- embeddings = model.encode(["find python json parser", "weather API tool"])
7
-
8
- # Search
9
- scores, indices = model.search(query_emb, index_emb, top_k=10)
10
- """
11
- from __future__ import annotations
12
-
13
- import json
14
- from pathlib import Path
15
-
16
- import numpy as np
17
-
18
-
19
- class LF4StaticEmbedding:
20
- """Native LF4 4-bit static embedding model.
21
-
22
- Weights are stored as packed 4-bit integers with per-block FP16 scales/zeros.
23
- Total model size: ~3.5 MB (vs 29 MB FP32).
24
- """
25
-
26
- def __init__(self, packed, scales, zeros, tokenizer_data, config):
27
- self.packed = packed # uint8 (vocab, dim/2)
28
- self.scales = scales # float16 (vocab, num_blocks)
29
- self.zeros = zeros # float16 (vocab, num_blocks)
30
- self.config = config
31
- self.vocab_size = config["vocab_size"]
32
- self.dim = config["embedding_dim"]
33
- self.block_size = config["block_size"]
34
- self._tokenizer_data = tokenizer_data
35
- self._tokenizer = None
36
-
37
- self._embedding_table = self._dequantize_all()
38
-
39
- def _dequantize_all(self) -> np.ndarray:
40
- """Dequantize full embedding table to FP32 for fast token lookup."""
41
- N = self.packed.shape[0]
42
- D = self.dim
43
- B = self.block_size
44
-
45
- low = (self.packed & 0x0F).astype(np.float32)
46
- high = ((self.packed >> 4) & 0x0F).astype(np.float32)
47
- D_padded = self.packed.shape[1] * 2
48
-
49
- unpacked = np.empty((N, D_padded), dtype=np.float32)
50
- unpacked[:, 0::2] = low
51
- unpacked[:, 1::2] = high
52
-
53
- num_blocks = D_padded // B
54
- blocked = unpacked.reshape(N, num_blocks, B)
55
- s = self.scales.astype(np.float32)[:, :, None]
56
- z = self.zeros.astype(np.float32)[:, :, None]
57
-
58
- return (blocked * s + z).reshape(N, D_padded)[:, :D]
59
-
60
- @property
61
- def tokenizer(self):
62
- if self._tokenizer is None:
63
- from tokenizers import Tokenizer
64
- self._tokenizer = Tokenizer.from_str(self._tokenizer_data)
65
- return self._tokenizer
66
-
67
- @classmethod
68
- def from_pretrained(cls, path_or_id: str) -> LF4StaticEmbedding:
69
- """Load model from local path or HuggingFace Hub."""
70
- p = Path(path_or_id)
71
- if p.is_dir():
72
- model_path = str(p / "model.safetensors")
73
- config_path = p / "config.json"
74
- tok_path = str(p / "tokenizer.json")
75
- else:
76
- from huggingface_hub import hf_hub_download
77
- model_path = hf_hub_download(path_or_id, "model.safetensors")
78
- config_path = Path(hf_hub_download(path_or_id, "config.json"))
79
- tok_path = hf_hub_download(path_or_id, "tokenizer.json")
80
-
81
- from safetensors.numpy import load_file
82
- tensors = load_file(model_path)
83
- config = json.loads(config_path.read_text())
84
-
85
- return cls(
86
- packed=tensors["embedding_packed"],
87
- scales=tensors["embedding_scales"],
88
- zeros=tensors["embedding_zeros"],
89
- tokenizer_data=Path(tok_path).read_text(encoding="utf-8"),
90
- config=config,
91
- )
92
-
93
- def encode(self, texts: str | list[str], normalize: bool = True) -> np.ndarray:
94
- """Encode texts to embeddings.
95
-
96
- Args:
97
- texts: single string or list of strings
98
- normalize: L2-normalize output embeddings (default True for cosine sim)
99
-
100
- Returns:
101
- np.ndarray of shape (N, dim)
102
- """
103
- if isinstance(texts, str):
104
- texts = [texts]
105
-
106
- embeddings = np.zeros((len(texts), self.dim), dtype=np.float32)
107
-
108
- for i, text in enumerate(texts):
109
- encoded = self.tokenizer.encode(text)
110
- token_ids = encoded.ids
111
-
112
- valid_ids = [tid for tid in token_ids if 0 <= tid < self.vocab_size]
113
- if valid_ids:
114
- token_embs = self._embedding_table[valid_ids]
115
- embeddings[i] = token_embs.mean(axis=0)
116
-
117
- if normalize:
118
- norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
119
- norms = np.where(norms == 0, 1.0, norms)
120
- embeddings = embeddings / norms
121
-
122
- return embeddings
123
-
124
- def search(
125
- self,
126
- queries: np.ndarray,
127
- index: np.ndarray,
128
- top_k: int = 10,
129
- ) -> tuple[np.ndarray, np.ndarray]:
130
- """Cosine similarity search.
131
-
132
- Args:
133
- queries: (Q, D) query embeddings
134
- index: (N, D) document embeddings
135
- top_k: number of results
136
-
137
- Returns:
138
- (scores, indices) arrays
139
- """
140
- queries = np.asarray(queries, dtype=np.float32)
141
- index = np.asarray(index, dtype=np.float32)
142
- if queries.ndim == 1:
143
- queries = queries[None, :]
144
-
145
- qn = queries / (np.linalg.norm(queries, axis=1, keepdims=True) + 1e-8)
146
- dn = index / (np.linalg.norm(index, axis=1, keepdims=True) + 1e-8)
147
-
148
- scores = qn @ dn.T
149
-
150
- if top_k >= scores.shape[1]:
151
- idx = np.argsort(-scores, axis=1)
152
- return np.take_along_axis(scores, idx, 1), idx
153
-
154
- idx = np.argpartition(-scores, top_k, axis=1)[:, :top_k]
155
- s = np.take_along_axis(scores, idx, 1)
156
- order = np.argsort(-s, axis=1)
157
- return np.take_along_axis(s, order, 1), np.take_along_axis(idx, order, 1)
158
-
159
- @property
160
- def model_size_mb(self) -> float:
161
- return (self.packed.nbytes + self.scales.nbytes + self.zeros.nbytes) / 1e6
162
-
163
- def __repr__(self):
164
- return (
165
- f"LF4StaticEmbedding(vocab={self.vocab_size}, dim={self.dim}, "
166
- f"bits=4, size={self.model_size_mb:.1f}MB, "
167
- f"block_size={self.block_size})"
168
- )
File without changes
File without changes
File without changes