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.
- {vortexa-0.3.1 → vortexa-0.3.2}/PKG-INFO +3 -7
- {vortexa-0.3.1 → vortexa-0.3.2}/README.md +2 -5
- {vortexa-0.3.1 → vortexa-0.3.2}/pyproject.toml +1 -2
- vortexa-0.3.2/src/vortexa/core/embedding.py +88 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/indexer.py +1 -1
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa.egg-info/PKG-INFO +3 -7
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa.egg-info/SOURCES.txt +0 -1
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa.egg-info/requires.txt +0 -1
- vortexa-0.3.1/src/vortexa/core/embedding.py +0 -180
- vortexa-0.3.1/src/vortexa/core/lf4_model.py +0 -168
- {vortexa-0.3.1 → vortexa-0.3.2}/LICENSE +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/setup.cfg +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/__init__.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/__init__.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/chunking.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/context_engine.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/graph.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/inference.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/language.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/lf4_v4_model.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/types.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/v4_embedder.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/core/vortex_score.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/interfaces/__init__.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/interfaces/cli.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/interfaces/mcp_server.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/interfaces/watcher.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/search/__init__.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/search/ranking.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/search/search.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/search/tokens.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/storage/__init__.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/storage/bm25.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/storage/vector_store.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa/storage/walker.py +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa.egg-info/dependency_links.txt +0 -0
- {vortexa-0.3.1 → vortexa-0.3.2}/src/vortexa.egg-info/entry_points.txt +0 -0
- {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.
|
|
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
|
|
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
|
|
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.
|
|
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
|
|
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.
|
|
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
|
|
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
|
|
|
@@ -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
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|