egemma2 0.3.1__py3-none-win_amd64.whl
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.
- egemma2/__init__.py +22 -0
- egemma2/__main__.py +17 -0
- egemma2/client.py +277 -0
- egemma2-0.3.1.data/scripts/egemma2.exe +0 -0
- egemma2-0.3.1.dist-info/METADATA +113 -0
- egemma2-0.3.1.dist-info/RECORD +10 -0
- egemma2-0.3.1.dist-info/WHEEL +4 -0
- egemma2-0.3.1.dist-info/licenses/LICENSE +202 -0
- egemma2-0.3.1.dist-info/licenses/NOTICE +12 -0
- egemma2-0.3.1.dist-info/sboms/embeddinggemma2.cyclonedx.json +6952 -0
egemma2/__init__.py
ADDED
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
"""EmbeddingGemma 2 in Rust: text, image, audio and video embeddings in one 768-d space, on CPU or CUDA.
|
|
2
|
+
|
|
3
|
+
`pip install egemma2` installs the `egemma2` command-line binary and this Python client for it:
|
|
4
|
+
|
|
5
|
+
from egemma2 import Egemma2
|
|
6
|
+
|
|
7
|
+
with Egemma2(precision="bf16") as model:
|
|
8
|
+
vectors = model.embed(["a query", "a document"], prompt_name="Document")
|
|
9
|
+
|
|
10
|
+
See https://github.com/Dalailalama/embeddinggemma2-rust for the documentation.
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
from .client import Egemma2, Egemma2Error, find_egemma2_bin, find_model_dir
|
|
14
|
+
|
|
15
|
+
try:
|
|
16
|
+
from importlib.metadata import PackageNotFoundError, version
|
|
17
|
+
|
|
18
|
+
__version__ = version("egemma2")
|
|
19
|
+
except (ImportError, PackageNotFoundError): # used from a source checkout
|
|
20
|
+
__version__ = "0+unknown"
|
|
21
|
+
|
|
22
|
+
__all__ = ["Egemma2", "Egemma2Error", "find_egemma2_bin", "find_model_dir", "__version__"]
|
egemma2/__main__.py
ADDED
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
"""`python -m egemma2 ...` runs the egemma2 binary installed with this package (same as `egemma2 ...`)."""
|
|
2
|
+
|
|
3
|
+
import subprocess
|
|
4
|
+
import sys
|
|
5
|
+
|
|
6
|
+
from .client import find_egemma2_bin
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
def main():
|
|
10
|
+
try:
|
|
11
|
+
sys.exit(subprocess.call([find_egemma2_bin(), *sys.argv[1:]]))
|
|
12
|
+
except KeyboardInterrupt:
|
|
13
|
+
sys.exit(130)
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
if __name__ == "__main__":
|
|
17
|
+
main()
|
egemma2/client.py
ADDED
|
@@ -0,0 +1,277 @@
|
|
|
1
|
+
"""Python client for `egemma2 serve`: the model stays loaded in one subprocess for all your requests.
|
|
2
|
+
|
|
3
|
+
Standard library only; numpy and Pillow are used when you pass their objects (or call `embed_numpy`), and
|
|
4
|
+
huggingface_hub, if installed, to download the model. `pip install egemma2` puts this module and the egemma2
|
|
5
|
+
binary it drives into the environment; the file also works on its own next to any egemma2 binary.
|
|
6
|
+
|
|
7
|
+
from egemma2 import Egemma2
|
|
8
|
+
|
|
9
|
+
with Egemma2(precision="bf16") as model: # CUDA if available, else the CPU
|
|
10
|
+
docs = model.embed(["first document", "second document"], prompt_name="Document")
|
|
11
|
+
query = model.embed("what is in the documents?", prompt_name="SearchQuery")[0]
|
|
12
|
+
pages = model.embed([Path("page1.png"), Path("page2.png")]) # files, by extension
|
|
13
|
+
mixed = model.embed([{"text": "a photo of", "image": pil_image}]) # one interleaved input
|
|
14
|
+
|
|
15
|
+
Inputs: a `str` is text. A `pathlib.Path` is a file chosen by its extension (image, .wav audio, or a video
|
|
16
|
+
file / frame directory). Images can also be Pillow images, HxWx3 uint8 numpy arrays or encoded bytes. A dict
|
|
17
|
+
such as `{"text": ..., "image": ...}` is ONE input whose parts follow the dict's order (like
|
|
18
|
+
sentence-transformers); `{"parts": [...]}` takes any sequence of parts. Audio given as data is
|
|
19
|
+
`(samples, sample_rate)`; a video can be a list of frames, with `{"video": frames, "fps": 30}`.
|
|
20
|
+
|
|
21
|
+
The server is started with `egemma2 [options] serve` and talks JSON lines over its stdin/stdout; see the
|
|
22
|
+
README's "Server mode" section. For an HTTP server instead, run `egemma2 serve --http 127.0.0.1:8080` and use
|
|
23
|
+
any OpenAI-compatible client.
|
|
24
|
+
"""
|
|
25
|
+
|
|
26
|
+
import base64
|
|
27
|
+
import json
|
|
28
|
+
import os
|
|
29
|
+
import shutil
|
|
30
|
+
import subprocess
|
|
31
|
+
import sys
|
|
32
|
+
import sysconfig
|
|
33
|
+
import threading
|
|
34
|
+
from array import array
|
|
35
|
+
from pathlib import Path
|
|
36
|
+
|
|
37
|
+
__all__ = ["Egemma2", "Egemma2Error", "find_egemma2_bin", "find_model_dir"]
|
|
38
|
+
|
|
39
|
+
MODEL_REPO = "google/embeddinggemma-2"
|
|
40
|
+
|
|
41
|
+
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".gif", ".bmp", ".tif", ".tiff"}
|
|
42
|
+
AUDIO_EXTS = {".wav", ".wave"}
|
|
43
|
+
VIDEO_EXTS = {".mp4", ".mov", ".mkv", ".webm", ".avi", ".m4v"}
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
class Egemma2Error(RuntimeError):
|
|
47
|
+
"""The server rejected a request, or stopped."""
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
def _b64(data):
|
|
51
|
+
return base64.b64encode(data).decode("ascii")
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def _image(v):
|
|
55
|
+
if isinstance(v, (str, os.PathLike)):
|
|
56
|
+
s = os.fspath(v)
|
|
57
|
+
return s if s.startswith("data:") else str(Path(s).resolve())
|
|
58
|
+
if isinstance(v, (bytes, bytearray, memoryview)):
|
|
59
|
+
return {"base64": _b64(bytes(v))}
|
|
60
|
+
if hasattr(v, "convert") and hasattr(v, "tobytes"): # Pillow image
|
|
61
|
+
rgb = v.convert("RGB")
|
|
62
|
+
return {"rgb": _b64(rgb.tobytes()), "width": rgb.width, "height": rgb.height}
|
|
63
|
+
shape = getattr(v, "shape", None)
|
|
64
|
+
if shape is not None and len(shape) == 3 and shape[2] == 3 and str(getattr(v, "dtype", "")) == "uint8":
|
|
65
|
+
import numpy as np
|
|
66
|
+
|
|
67
|
+
return {"rgb": _b64(np.ascontiguousarray(v).tobytes()), "width": int(shape[1]), "height": int(shape[0])}
|
|
68
|
+
raise TypeError(f"unsupported image value: {type(v).__name__}")
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
def _audio(v):
|
|
72
|
+
if isinstance(v, (str, os.PathLike)):
|
|
73
|
+
return str(Path(v).resolve())
|
|
74
|
+
if isinstance(v, (bytes, bytearray, memoryview)): # an encoded .wav file
|
|
75
|
+
return {"base64": _b64(bytes(v))}
|
|
76
|
+
if isinstance(v, tuple) and len(v) == 2: # (samples, sample_rate)
|
|
77
|
+
samples, rate = v
|
|
78
|
+
if hasattr(samples, "astype"):
|
|
79
|
+
raw = samples.astype("<f4").tobytes()
|
|
80
|
+
else:
|
|
81
|
+
a = array("f", samples)
|
|
82
|
+
if sys.byteorder == "big":
|
|
83
|
+
a.byteswap()
|
|
84
|
+
raw = a.tobytes()
|
|
85
|
+
return {"pcm_f32": _b64(raw), "sample_rate": int(rate)}
|
|
86
|
+
raise TypeError(f"unsupported audio value: {type(v).__name__} (use a path, wav bytes or (samples, sample_rate))")
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
def _video(v):
|
|
90
|
+
if isinstance(v, (str, os.PathLike)):
|
|
91
|
+
return str(Path(v).resolve())
|
|
92
|
+
if isinstance(v, (list, tuple)):
|
|
93
|
+
return [_image(f) for f in v]
|
|
94
|
+
raise TypeError(f"unsupported video value: {type(v).__name__} (use a path or a list of frames)")
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
def _path_item(p):
|
|
98
|
+
p = Path(p)
|
|
99
|
+
ext = p.suffix.lower()
|
|
100
|
+
if p.is_dir() or ext in VIDEO_EXTS:
|
|
101
|
+
return {"video": _video(p)}
|
|
102
|
+
if ext in AUDIO_EXTS:
|
|
103
|
+
return {"audio": _audio(p)}
|
|
104
|
+
if ext in IMAGE_EXTS:
|
|
105
|
+
return {"image": _image(p)}
|
|
106
|
+
raise ValueError(f"cannot tell what {p} is from its extension; pass {{'image': ...}}, {{'audio': ...}} or {{'video': ...}}")
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
_CONVERT = {"image": _image, "image_url": _image, "audio": _audio, "video": _video}
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
def _item(x):
|
|
113
|
+
if isinstance(x, str):
|
|
114
|
+
return x
|
|
115
|
+
if isinstance(x, os.PathLike):
|
|
116
|
+
return _path_item(x)
|
|
117
|
+
if isinstance(x, dict):
|
|
118
|
+
out = {}
|
|
119
|
+
for k, v in x.items(): # dicts keep insertion order, which is the parts' order
|
|
120
|
+
if k in _CONVERT:
|
|
121
|
+
out[k] = _CONVERT[k](v)
|
|
122
|
+
elif k in ("parts", "content"):
|
|
123
|
+
out[k] = [_item(p) for p in v]
|
|
124
|
+
else:
|
|
125
|
+
out[k] = v
|
|
126
|
+
return out
|
|
127
|
+
try:
|
|
128
|
+
return {"image": _image(x)}
|
|
129
|
+
except TypeError:
|
|
130
|
+
raise TypeError(f"unsupported input: {type(x).__name__}") from None
|
|
131
|
+
|
|
132
|
+
|
|
133
|
+
def _floats(v):
|
|
134
|
+
if isinstance(v, str):
|
|
135
|
+
a = array("f")
|
|
136
|
+
a.frombytes(base64.b64decode(v))
|
|
137
|
+
if sys.byteorder == "big":
|
|
138
|
+
a.byteswap()
|
|
139
|
+
return a
|
|
140
|
+
return v
|
|
141
|
+
|
|
142
|
+
|
|
143
|
+
def find_egemma2_bin():
|
|
144
|
+
"""The egemma2 binary: $EGEMMA2_EXE, the one a `pip install egemma2` put next to this Python, or the PATH."""
|
|
145
|
+
env = os.environ.get("EGEMMA2_EXE")
|
|
146
|
+
if env:
|
|
147
|
+
return env
|
|
148
|
+
name = "egemma2" + (sysconfig.get_config_var("EXE") or (".exe" if os.name == "nt" else ""))
|
|
149
|
+
dirs = [sysconfig.get_path("scripts")]
|
|
150
|
+
user_scheme = getattr(sysconfig, "get_preferred_scheme", lambda _: "nt_user" if os.name == "nt" else "posix_user")("user")
|
|
151
|
+
try:
|
|
152
|
+
dirs.append(sysconfig.get_path("scripts", scheme=user_scheme)) # pip install --user
|
|
153
|
+
except KeyError:
|
|
154
|
+
pass
|
|
155
|
+
dirs.append(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "bin")) # pip install --target
|
|
156
|
+
for d in dirs:
|
|
157
|
+
path = os.path.join(d, name)
|
|
158
|
+
if os.path.isfile(path):
|
|
159
|
+
return path
|
|
160
|
+
found = shutil.which("egemma2")
|
|
161
|
+
if found:
|
|
162
|
+
return found
|
|
163
|
+
raise FileNotFoundError("egemma2 not found: `pip install egemma2`, pass exe=..., set EGEMMA2_EXE or put it on PATH")
|
|
164
|
+
|
|
165
|
+
|
|
166
|
+
def find_model_dir():
|
|
167
|
+
"""The model directory: $EGEMMA2_MODEL_DIR, ./models/embeddinggemma-2, or a download of google/embeddinggemma-2
|
|
168
|
+
into the Hugging Face cache (needs huggingface_hub: `pip install "egemma2[hub]"`; about 1.5 GB, once)."""
|
|
169
|
+
env = os.environ.get("EGEMMA2_MODEL_DIR")
|
|
170
|
+
if env:
|
|
171
|
+
return env
|
|
172
|
+
local = Path("models") / "embeddinggemma-2"
|
|
173
|
+
if (local / "model.safetensors").is_file():
|
|
174
|
+
return str(local.resolve())
|
|
175
|
+
try:
|
|
176
|
+
from huggingface_hub import snapshot_download
|
|
177
|
+
except ImportError:
|
|
178
|
+
raise FileNotFoundError(
|
|
179
|
+
"no model found: pass model_dir=..., set EGEMMA2_MODEL_DIR, install huggingface_hub to download it "
|
|
180
|
+
f'automatically (pip install "egemma2[hub]"), or run: hf download {MODEL_REPO} --local-dir models/embeddinggemma-2'
|
|
181
|
+
) from None
|
|
182
|
+
return snapshot_download(MODEL_REPO, allow_patterns=["*.json", "*.jinja", "model.safetensors"])
|
|
183
|
+
|
|
184
|
+
|
|
185
|
+
class Egemma2:
|
|
186
|
+
"""One `egemma2 serve` subprocess. Thread-safe: concurrent calls are sent one at a time."""
|
|
187
|
+
|
|
188
|
+
def __init__(self, exe=None, model_dir=None, device="auto", precision=None, max_tokens=None, args=(), stderr=None, cwd=None):
|
|
189
|
+
"""
|
|
190
|
+
exe: path to the egemma2 binary (default: see `find_egemma2_bin`).
|
|
191
|
+
model_dir: the model directory (default: see `find_model_dir`).
|
|
192
|
+
device: auto, cpu, cuda or cuda:N. precision: None (float32), "tf32" or "bf16" (GPU only).
|
|
193
|
+
max_tokens: cut texts longer than this (default: no cut, like sentence-transformers).
|
|
194
|
+
args: extra options for `egemma2 serve` (e.g. ["--prompt-name", "Document"]).
|
|
195
|
+
stderr: where the server's log goes (default: this process's stderr).
|
|
196
|
+
"""
|
|
197
|
+
exe = os.fspath(exe) if exe is not None else find_egemma2_bin()
|
|
198
|
+
# A relative path such as target/release/egemma2 is not searched for by Windows' CreateProcess.
|
|
199
|
+
exe = os.path.abspath(exe) if os.path.isfile(exe) else (shutil.which(exe) or exe)
|
|
200
|
+
model_dir = os.path.abspath(os.fspath(model_dir)) if model_dir is not None else find_model_dir()
|
|
201
|
+
cmd = [exe, "--device", device, "--model-dir", model_dir]
|
|
202
|
+
if precision in ("bf16", "tf32"):
|
|
203
|
+
cmd.append("--" + precision)
|
|
204
|
+
elif precision not in (None, "fp32", "f32", "float32"):
|
|
205
|
+
raise ValueError(f"precision must be None, 'tf32' or 'bf16' (got {precision!r})")
|
|
206
|
+
if max_tokens is not None:
|
|
207
|
+
cmd += ["--max-tokens", str(int(max_tokens))]
|
|
208
|
+
cmd += ["serve", *args]
|
|
209
|
+
self._proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=stderr, cwd=cwd)
|
|
210
|
+
self._lock = threading.Lock()
|
|
211
|
+
self._next_id = 0
|
|
212
|
+
|
|
213
|
+
def embed(self, inputs, *, prompt_name=None, prompt=None, dimensions=None):
|
|
214
|
+
"""Embed one input or a list of inputs; returns one list of floats (L2-normalised) per input.
|
|
215
|
+
|
|
216
|
+
prompt_name: a sentence-transformers prompt such as "SearchQuery" or "Document" (or a literal
|
|
217
|
+
`prompt` string), placed before the whole input. dimensions: Matryoshka size (768, 512, 256, 128).
|
|
218
|
+
"""
|
|
219
|
+
single = isinstance(inputs, (str, dict, os.PathLike)) or not isinstance(inputs, (list, tuple))
|
|
220
|
+
items = [_item(inputs)] if single else [_item(x) for x in inputs]
|
|
221
|
+
if not items:
|
|
222
|
+
return []
|
|
223
|
+
req = {"input": items, "encoding_format": "base64"}
|
|
224
|
+
if prompt_name is not None:
|
|
225
|
+
req["prompt_name"] = prompt_name
|
|
226
|
+
if prompt is not None:
|
|
227
|
+
req["prompt"] = prompt
|
|
228
|
+
if dimensions is not None:
|
|
229
|
+
req["dimensions"] = int(dimensions)
|
|
230
|
+
resp = self._call(req)
|
|
231
|
+
return [list(_floats(d["embedding"])) for d in sorted(resp["data"], key=lambda d: d["index"])]
|
|
232
|
+
|
|
233
|
+
def embed_numpy(self, inputs, **kwargs):
|
|
234
|
+
"""Like `embed`, as a float32 numpy array of shape (n, dim)."""
|
|
235
|
+
import numpy as np
|
|
236
|
+
|
|
237
|
+
return np.asarray(self.embed(inputs, **kwargs), dtype=np.float32)
|
|
238
|
+
|
|
239
|
+
def _call(self, req):
|
|
240
|
+
with self._lock:
|
|
241
|
+
if self._proc.poll() is not None:
|
|
242
|
+
raise Egemma2Error(f"the egemma2 server has exited (code {self._proc.returncode})")
|
|
243
|
+
self._next_id += 1
|
|
244
|
+
req["id"] = self._next_id
|
|
245
|
+
try:
|
|
246
|
+
self._proc.stdin.write(json.dumps(req).encode("utf-8") + b"\n")
|
|
247
|
+
self._proc.stdin.flush()
|
|
248
|
+
line = self._proc.stdout.readline()
|
|
249
|
+
except (BrokenPipeError, OSError) as e:
|
|
250
|
+
raise Egemma2Error(f"the egemma2 server stopped: {e}") from None
|
|
251
|
+
if not line:
|
|
252
|
+
raise Egemma2Error(f"the egemma2 server stopped (exit code {self._proc.wait()}); see its log above")
|
|
253
|
+
resp = json.loads(line)
|
|
254
|
+
if "error" in resp:
|
|
255
|
+
raise Egemma2Error(resp["error"]["message"])
|
|
256
|
+
return resp
|
|
257
|
+
|
|
258
|
+
def close(self):
|
|
259
|
+
"""Stop the server (it exits when its stdin closes)."""
|
|
260
|
+
if self._proc.poll() is None:
|
|
261
|
+
try:
|
|
262
|
+
self._proc.stdin.close()
|
|
263
|
+
except OSError:
|
|
264
|
+
pass
|
|
265
|
+
try:
|
|
266
|
+
self._proc.wait(timeout=30)
|
|
267
|
+
except subprocess.TimeoutExpired:
|
|
268
|
+
self._proc.kill()
|
|
269
|
+
self._proc.wait()
|
|
270
|
+
if self._proc.stdout is not None:
|
|
271
|
+
self._proc.stdout.close()
|
|
272
|
+
|
|
273
|
+
def __enter__(self):
|
|
274
|
+
return self
|
|
275
|
+
|
|
276
|
+
def __exit__(self, *exc):
|
|
277
|
+
self.close()
|
|
Binary file
|
|
@@ -0,0 +1,113 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: egemma2
|
|
3
|
+
Version: 0.3.1
|
|
4
|
+
Classifier: Development Status :: 4 - Beta
|
|
5
|
+
Classifier: Environment :: GPU :: NVIDIA CUDA
|
|
6
|
+
Classifier: Intended Audience :: Developers
|
|
7
|
+
Classifier: Intended Audience :: Science/Research
|
|
8
|
+
Classifier: Operating System :: Microsoft :: Windows
|
|
9
|
+
Classifier: Operating System :: POSIX :: Linux
|
|
10
|
+
Classifier: Programming Language :: Python :: 3
|
|
11
|
+
Classifier: Programming Language :: Rust
|
|
12
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
13
|
+
Classifier: Topic :: Text Processing :: Indexing
|
|
14
|
+
Requires-Dist: huggingface-hub>=0.23 ; extra == 'hub'
|
|
15
|
+
Requires-Dist: numpy ; extra == 'numpy'
|
|
16
|
+
Provides-Extra: hub
|
|
17
|
+
Provides-Extra: numpy
|
|
18
|
+
License-File: LICENSE
|
|
19
|
+
License-File: NOTICE
|
|
20
|
+
Summary: EmbeddingGemma 2 in Rust: text, image, audio and video embeddings on CPU or CUDA, with a Python client and an OpenAI-compatible server
|
|
21
|
+
Keywords: embeddings,embeddinggemma,gemma,multimodal,semantic-search,rag,cuda,rust
|
|
22
|
+
Author: DawgOG
|
|
23
|
+
License-Expression: Apache-2.0
|
|
24
|
+
Requires-Python: >=3.8
|
|
25
|
+
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
|
|
26
|
+
Project-URL: Documentation, https://github.com/Dalailalama/embeddinggemma2-rust#server-mode-keep-the-model-loaded
|
|
27
|
+
Project-URL: Funding, https://buymeacoffee.com/dawgog
|
|
28
|
+
Project-URL: Homepage, https://github.com/Dalailalama/embeddinggemma2-rust
|
|
29
|
+
Project-URL: Hugging Face, https://huggingface.co/DawgOG/embeddinggemma2-rust
|
|
30
|
+
|
|
31
|
+
# egemma2: EmbeddingGemma 2 in Rust, for Python
|
|
32
|
+
|
|
33
|
+
An unofficial, from-scratch Rust implementation of Google's
|
|
34
|
+
[EmbeddingGemma 2](https://huggingface.co/google/embeddinggemma-2): text, images, audio and video
|
|
35
|
+
embedded into one 768-d space (Matryoshka cuts to 512 / 256 / 128), on the CPU or an NVIDIA GPU.
|
|
36
|
+
This package installs the `egemma2` binary plus a small Python client that keeps it loaded, so
|
|
37
|
+
Python code gets the speed without PyTorch:
|
|
38
|
+
|
|
39
|
+
| Workload (RTX A3000 laptop GPU, bf16) | PyTorch + sentence-transformers | egemma2 |
|
|
40
|
+
|---|---|---|
|
|
41
|
+
| short query | 52.3 ms | **2.5 ms** |
|
|
42
|
+
| 2,208-token document | 136 ms | **47.6 ms** |
|
|
43
|
+
| 16 images per call, per image | 99.9 ms | **51.4 ms** |
|
|
44
|
+
| 9.3 s audio clip | 91.0 ms | **18.9 ms** |
|
|
45
|
+
|
|
46
|
+
Embeddings match the PyTorch reference: up to 1e-6 apart in float32, and within cosine 0.99996 in
|
|
47
|
+
bf16 mode. Texts are not truncated, as in sentence-transformers.
|
|
48
|
+
|
|
49
|
+
## Install
|
|
50
|
+
|
|
51
|
+
```bash
|
|
52
|
+
pip install "egemma2[hub]" # with huggingface_hub, to download the model on first use (~1.5 GB, once)
|
|
53
|
+
pip install egemma2 # or just the binary and the client, for a model you already have
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
Wheels: Windows and Linux x64, any Python 3.8+. The GPU path needs an NVIDIA driver and the CUDA 13
|
|
57
|
+
runtime libraries (cuBLAS, NVRTC) on `PATH`, for example from the CUDA Toolkit; bf16 needs compute
|
|
58
|
+
capability 8.0+ (Ampere or newer). Without them egemma2 runs on the CPU.
|
|
59
|
+
|
|
60
|
+
## Use
|
|
61
|
+
|
|
62
|
+
```python
|
|
63
|
+
from pathlib import Path
|
|
64
|
+
from egemma2 import Egemma2
|
|
65
|
+
|
|
66
|
+
with Egemma2(precision="bf16") as model: # device="auto": CUDA if available, else the CPU
|
|
67
|
+
docs = model.embed(["Mars is red.", "Oceans are blue."], prompt_name="Document")
|
|
68
|
+
query = model.embed("which planet is red?", prompt_name="SearchQuery")[0]
|
|
69
|
+
scores = [sum(q * d for q, d in zip(query, doc)) for doc in docs] # cosine similarity
|
|
70
|
+
|
|
71
|
+
pages = model.embed([Path("page1.png"), Path("page2.png")]) # files, by extension
|
|
72
|
+
mixed = model.embed([{"text": "Figure 3:", "image": pil_image}]) # one interleaved input
|
|
73
|
+
matrix = model.embed_numpy(texts, dimensions=256) # (n, 256) float32
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
- **Inputs:** a `str` is text, a `Path` is a file (image, `.wav` audio, video file or frame
|
|
77
|
+
directory). Images can also be Pillow images, HxWx3 uint8 numpy arrays or encoded bytes; audio can
|
|
78
|
+
be `(samples, sample_rate)`. A dict such as `{"text": ..., "image": ...}` is **one** input whose
|
|
79
|
+
parts follow the dict's order, like sentence-transformers' dict inputs.
|
|
80
|
+
- **Prompts:** `prompt_name` takes the model's sentence-transformers prompts (`SearchQuery`,
|
|
81
|
+
`Document`, `QuestionAnswering`, `CodeRetrieval`, ...).
|
|
82
|
+
- **Model directory:** `Egemma2(model_dir=...)`, else `$EGEMMA2_MODEL_DIR`, else
|
|
83
|
+
`./models/embeddinggemma-2`, else a download of `google/embeddinggemma-2` via huggingface_hub.
|
|
84
|
+
- **Other options:** `device="cpu" | "cuda" | "cuda:N"`, `precision=None | "tf32" | "bf16"`,
|
|
85
|
+
`max_tokens=N` to cut long texts.
|
|
86
|
+
|
|
87
|
+
The client starts `egemma2 serve` once and sends it JSON lines, so the model loads once; a short
|
|
88
|
+
query takes 2.6 ms round trip. Requests from several threads are queued safely.
|
|
89
|
+
|
|
90
|
+
## Command line and HTTP server
|
|
91
|
+
|
|
92
|
+
The same install gives you the `egemma2` command (also `python -m egemma2`):
|
|
93
|
+
|
|
94
|
+
```bash
|
|
95
|
+
egemma2 --model-dir <dir> embed --prompt-name SearchQuery --text "which planet is red?"
|
|
96
|
+
egemma2 --model-dir <dir> --bf16 serve --http 127.0.0.1:8080 # OpenAI-compatible /v1/embeddings
|
|
97
|
+
```
|
|
98
|
+
|
|
99
|
+
Any OpenAI-style client (the `openai` SDK, LangChain, LlamaIndex) can use the HTTP server:
|
|
100
|
+
|
|
101
|
+
```python
|
|
102
|
+
from openai import OpenAI
|
|
103
|
+
client = OpenAI(base_url="http://127.0.0.1:8080/v1", api_key="unused")
|
|
104
|
+
r = client.embeddings.create(model="embeddinggemma-2", input=["a document"], extra_body={"prompt_name": "Document"})
|
|
105
|
+
```
|
|
106
|
+
|
|
107
|
+
## More
|
|
108
|
+
|
|
109
|
+
Documentation, benchmarks, the parity tests and the source:
|
|
110
|
+
<https://github.com/Dalailalama/embeddinggemma2-rust>. Not affiliated with or endorsed by Google;
|
|
111
|
+
the model weights are Google's (Apache-2.0). Code: Apache-2.0. If it saves you time, you can
|
|
112
|
+
[buy me a coffee](https://buymeacoffee.com/dawgog).
|
|
113
|
+
|
|
@@ -0,0 +1,10 @@
|
|
|
1
|
+
egemma2/__init__.py,sha256=-d8VtBroazNGKD3zkjQcqw-6QZwr2DOFdKcPkuL7Ycs,829
|
|
2
|
+
egemma2/__main__.py,sha256=5Nm17EBUcRCaAHKEAHS5lzpa94fPpg3UGpoPEoMBLFc,361
|
|
3
|
+
egemma2/client.py,sha256=5FxCehhMDQ1ZNR0c_Oa2WRO0BG7_NcMyBvrFJvry2hc,11818
|
|
4
|
+
egemma2-0.3.1.data/scripts/egemma2.exe,sha256=1G3NgzZb4z52P_O8tYYXu_EnutYUs4Dafdl3hg22Gk8,8970752
|
|
5
|
+
egemma2-0.3.1.dist-info/METADATA,sha256=hLVipXWfaDfptV6AJWJBrxBkgU1IJ1hs80JJQfB-cB4,5475
|
|
6
|
+
egemma2-0.3.1.dist-info/WHEEL,sha256=8Aej0W0a6Cz6apA3IzJrTnxLRVLAt-w0Oh8SA3Con_c,94
|
|
7
|
+
egemma2-0.3.1.dist-info/licenses/LICENSE,sha256=Pd-b5cKP4n2tFDpdx27qJSIq0d1ok0oEcGTlbtL6QMU,11560
|
|
8
|
+
egemma2-0.3.1.dist-info/licenses/NOTICE,sha256=DGVNMScQwtnm-xdyZISnvm4Tac1Mx81n3Pk1o6V3Rmk,734
|
|
9
|
+
egemma2-0.3.1.dist-info/sboms/embeddinggemma2.cyclonedx.json,sha256=QTAjgn4KN3SwjXZ21bo-qXhWtKLDDLjdL_wdMEGFHWU,222338
|
|
10
|
+
egemma2-0.3.1.dist-info/RECORD,,
|
|
@@ -0,0 +1,202 @@
|
|
|
1
|
+
|
|
2
|
+
Apache License
|
|
3
|
+
Version 2.0, January 2004
|
|
4
|
+
http://www.apache.org/licenses/
|
|
5
|
+
|
|
6
|
+
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
|
7
|
+
|
|
8
|
+
1. Definitions.
|
|
9
|
+
|
|
10
|
+
"License" shall mean the terms and conditions for use, reproduction,
|
|
11
|
+
and distribution as defined by Sections 1 through 9 of this document.
|
|
12
|
+
|
|
13
|
+
"Licensor" shall mean the copyright owner or entity authorized by
|
|
14
|
+
the copyright owner that is granting the License.
|
|
15
|
+
|
|
16
|
+
"Legal Entity" shall mean the union of the acting entity and all
|
|
17
|
+
other entities that control, are controlled by, or are under common
|
|
18
|
+
control with that entity. For the purposes of this definition,
|
|
19
|
+
"control" means (i) the power, direct or indirect, to cause the
|
|
20
|
+
direction or management of such entity, whether by contract or
|
|
21
|
+
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
|
22
|
+
outstanding shares, or (iii) beneficial ownership of such entity.
|
|
23
|
+
|
|
24
|
+
"You" (or "Your") shall mean an individual or Legal Entity
|
|
25
|
+
exercising permissions granted by this License.
|
|
26
|
+
|
|
27
|
+
"Source" form shall mean the preferred form for making modifications,
|
|
28
|
+
including but not limited to software source code, documentation
|
|
29
|
+
source, and configuration files.
|
|
30
|
+
|
|
31
|
+
"Object" form shall mean any form resulting from mechanical
|
|
32
|
+
transformation or translation of a Source form, including but
|
|
33
|
+
not limited to compiled object code, generated documentation,
|
|
34
|
+
and conversions to other media types.
|
|
35
|
+
|
|
36
|
+
"Work" shall mean the work of authorship, whether in Source or
|
|
37
|
+
Object form, made available under the License, as indicated by a
|
|
38
|
+
copyright notice that is included in or attached to the work
|
|
39
|
+
(an example is provided in the Appendix below).
|
|
40
|
+
|
|
41
|
+
"Derivative Works" shall mean any work, whether in Source or Object
|
|
42
|
+
form, that is based on (or derived from) the Work and for which the
|
|
43
|
+
editorial revisions, annotations, elaborations, or other modifications
|
|
44
|
+
represent, as a whole, an original work of authorship. For the purposes
|
|
45
|
+
of this License, Derivative Works shall not include works that remain
|
|
46
|
+
separable from, or merely link (or bind by name) to the interfaces of,
|
|
47
|
+
the Work and Derivative Works thereof.
|
|
48
|
+
|
|
49
|
+
"Contribution" shall mean any work of authorship, including
|
|
50
|
+
the original version of the Work and any modifications or additions
|
|
51
|
+
to that Work or Derivative Works thereof, that is intentionally
|
|
52
|
+
submitted to Licensor for inclusion in the Work by the copyright owner
|
|
53
|
+
or by an individual or Legal Entity authorized to submit on behalf of
|
|
54
|
+
the copyright owner. For the purposes of this definition, "submitted"
|
|
55
|
+
means any form of electronic, verbal, or written communication sent
|
|
56
|
+
to the Licensor or its representatives, including but not limited to
|
|
57
|
+
communication on electronic mailing lists, source code control systems,
|
|
58
|
+
and issue tracking systems that are managed by, or on behalf of, the
|
|
59
|
+
Licensor for the purpose of discussing and improving the Work, but
|
|
60
|
+
excluding communication that is conspicuously marked or otherwise
|
|
61
|
+
designated in writing by the copyright owner as "Not a Contribution."
|
|
62
|
+
|
|
63
|
+
"Contributor" shall mean Licensor and any individual or Legal Entity
|
|
64
|
+
on behalf of whom a Contribution has been received by Licensor and
|
|
65
|
+
subsequently incorporated within the Work.
|
|
66
|
+
|
|
67
|
+
2. Grant of Copyright License. Subject to the terms and conditions of
|
|
68
|
+
this License, each Contributor hereby grants to You a perpetual,
|
|
69
|
+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
|
70
|
+
copyright license to reproduce, prepare Derivative Works of,
|
|
71
|
+
publicly display, publicly perform, sublicense, and distribute the
|
|
72
|
+
Work and such Derivative Works in Source or Object form.
|
|
73
|
+
|
|
74
|
+
3. Grant of Patent License. Subject to the terms and conditions of
|
|
75
|
+
this License, each Contributor hereby grants to You a perpetual,
|
|
76
|
+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
|
77
|
+
(except as stated in this section) patent license to make, have made,
|
|
78
|
+
use, offer to sell, sell, import, and otherwise transfer the Work,
|
|
79
|
+
where such license applies only to those patent claims licensable
|
|
80
|
+
by such Contributor that are necessarily infringed by their
|
|
81
|
+
Contribution(s) alone or by combination of their Contribution(s)
|
|
82
|
+
with the Work to which such Contribution(s) was submitted. If You
|
|
83
|
+
institute patent litigation against any entity (including a
|
|
84
|
+
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
|
85
|
+
or a Contribution incorporated within the Work constitutes direct
|
|
86
|
+
or contributory patent infringement, then any patent licenses
|
|
87
|
+
granted to You under this License for that Work shall terminate
|
|
88
|
+
as of the date such litigation is filed.
|
|
89
|
+
|
|
90
|
+
4. Redistribution. You may reproduce and distribute copies of the
|
|
91
|
+
Work or Derivative Works thereof in any medium, with or without
|
|
92
|
+
modifications, and in Source or Object form, provided that You
|
|
93
|
+
meet the following conditions:
|
|
94
|
+
|
|
95
|
+
(a) You must give any other recipients of the Work or
|
|
96
|
+
Derivative Works a copy of this License; and
|
|
97
|
+
|
|
98
|
+
(b) You must cause any modified files to carry prominent notices
|
|
99
|
+
stating that You changed the files; and
|
|
100
|
+
|
|
101
|
+
(c) You must retain, in the Source form of any Derivative Works
|
|
102
|
+
that You distribute, all copyright, patent, trademark, and
|
|
103
|
+
attribution notices from the Source form of the Work,
|
|
104
|
+
excluding those notices that do not pertain to any part of
|
|
105
|
+
the Derivative Works; and
|
|
106
|
+
|
|
107
|
+
(d) If the Work includes a "NOTICE" text file as part of its
|
|
108
|
+
distribution, then any Derivative Works that You distribute must
|
|
109
|
+
include a readable copy of the attribution notices contained
|
|
110
|
+
within such NOTICE file, excluding those notices that do not
|
|
111
|
+
pertain to any part of the Derivative Works, in at least one
|
|
112
|
+
of the following places: within a NOTICE text file distributed
|
|
113
|
+
as part of the Derivative Works; within the Source form or
|
|
114
|
+
documentation, if provided along with the Derivative Works; or,
|
|
115
|
+
within a display generated by the Derivative Works, if and
|
|
116
|
+
wherever such third-party notices normally appear. The contents
|
|
117
|
+
of the NOTICE file are for informational purposes only and
|
|
118
|
+
do not modify the License. You may add Your own attribution
|
|
119
|
+
notices within Derivative Works that You distribute, alongside
|
|
120
|
+
or as an addendum to the NOTICE text from the Work, provided
|
|
121
|
+
that such additional attribution notices cannot be construed
|
|
122
|
+
as modifying the License.
|
|
123
|
+
|
|
124
|
+
You may add Your own copyright statement to Your modifications and
|
|
125
|
+
may provide additional or different license terms and conditions
|
|
126
|
+
for use, reproduction, or distribution of Your modifications, or
|
|
127
|
+
for any such Derivative Works as a whole, provided Your use,
|
|
128
|
+
reproduction, and distribution of the Work otherwise complies with
|
|
129
|
+
the conditions stated in this License.
|
|
130
|
+
|
|
131
|
+
5. Submission of Contributions. Unless You explicitly state otherwise,
|
|
132
|
+
any Contribution intentionally submitted for inclusion in the Work
|
|
133
|
+
by You to the Licensor shall be under the terms and conditions of
|
|
134
|
+
this License, without any additional terms or conditions.
|
|
135
|
+
Notwithstanding the above, nothing herein shall supersede or modify
|
|
136
|
+
the terms of any separate license agreement you may have executed
|
|
137
|
+
with Licensor regarding such Contributions.
|
|
138
|
+
|
|
139
|
+
6. Trademarks. This License does not grant permission to use the trade
|
|
140
|
+
names, trademarks, service marks, or product names of the Licensor,
|
|
141
|
+
except as required for reasonable and customary use in describing the
|
|
142
|
+
origin of the Work and reproducing the content of the NOTICE file.
|
|
143
|
+
|
|
144
|
+
7. Disclaimer of Warranty. Unless required by applicable law or
|
|
145
|
+
agreed to in writing, Licensor provides the Work (and each
|
|
146
|
+
Contributor provides its Contributions) on an "AS IS" BASIS,
|
|
147
|
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
|
148
|
+
implied, including, without limitation, any warranties or conditions
|
|
149
|
+
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
|
150
|
+
PARTICULAR PURPOSE. You are solely responsible for determining the
|
|
151
|
+
appropriateness of using or redistributing the Work and assume any
|
|
152
|
+
risks associated with Your exercise of permissions under this License.
|
|
153
|
+
|
|
154
|
+
8. Limitation of Liability. In no event and under no legal theory,
|
|
155
|
+
whether in tort (including negligence), contract, or otherwise,
|
|
156
|
+
unless required by applicable law (such as deliberate and grossly
|
|
157
|
+
negligent acts) or agreed to in writing, shall any Contributor be
|
|
158
|
+
liable to You for damages, including any direct, indirect, special,
|
|
159
|
+
incidental, or consequential damages of any character arising as a
|
|
160
|
+
result of this License or out of the use or inability to use the
|
|
161
|
+
Work (including but not limited to damages for loss of goodwill,
|
|
162
|
+
work stoppage, computer failure or malfunction, or any and all
|
|
163
|
+
other commercial damages or losses), even if such Contributor
|
|
164
|
+
has been advised of the possibility of such damages.
|
|
165
|
+
|
|
166
|
+
9. Accepting Warranty or Additional Liability. While redistributing
|
|
167
|
+
the Work or Derivative Works thereof, You may choose to offer,
|
|
168
|
+
and charge a fee for, acceptance of support, warranty, indemnity,
|
|
169
|
+
or other liability obligations and/or rights consistent with this
|
|
170
|
+
License. However, in accepting such obligations, You may act only
|
|
171
|
+
on Your own behalf and on Your sole responsibility, not on behalf
|
|
172
|
+
of any other Contributor, and only if You agree to indemnify,
|
|
173
|
+
defend, and hold each Contributor harmless for any liability
|
|
174
|
+
incurred by, or claims asserted against, such Contributor by reason
|
|
175
|
+
of your accepting any such warranty or additional liability.
|
|
176
|
+
|
|
177
|
+
END OF TERMS AND CONDITIONS
|
|
178
|
+
|
|
179
|
+
APPENDIX: How to apply the Apache License to your work.
|
|
180
|
+
|
|
181
|
+
To apply the Apache License to your work, attach the following
|
|
182
|
+
boilerplate notice, with the fields enclosed by brackets "[]"
|
|
183
|
+
replaced with your own identifying information. (Don't include
|
|
184
|
+
the brackets!) The text should be enclosed in the appropriate
|
|
185
|
+
comment syntax for the file format. We also recommend that a
|
|
186
|
+
file or class name and description of purpose be included on the
|
|
187
|
+
same "printed page" as the copyright notice for easier
|
|
188
|
+
identification within third-party archives.
|
|
189
|
+
|
|
190
|
+
Copyright [yyyy] [name of copyright owner]
|
|
191
|
+
|
|
192
|
+
Licensed under the Apache License, Version 2.0 (the "License");
|
|
193
|
+
you may not use this file except in compliance with the License.
|
|
194
|
+
You may obtain a copy of the License at
|
|
195
|
+
|
|
196
|
+
http://www.apache.org/licenses/LICENSE-2.0
|
|
197
|
+
|
|
198
|
+
Unless required by applicable law or agreed to in writing, software
|
|
199
|
+
distributed under the License is distributed on an "AS IS" BASIS,
|
|
200
|
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
201
|
+
See the License for the specific language governing permissions and
|
|
202
|
+
limitations under the License.
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
egemma2: EmbeddingGemma 2 in Rust
|
|
2
|
+
Copyright 2026 DawgOG (https://huggingface.co/DawgOG, https://github.com/Dalailalama)
|
|
3
|
+
|
|
4
|
+
This product is licensed under the Apache License, Version 2.0 (see LICENSE).
|
|
5
|
+
|
|
6
|
+
It is an unofficial, independent implementation of Google's EmbeddingGemma 2 model
|
|
7
|
+
(https://huggingface.co/google/embeddinggemma-2) and is not affiliated with, sponsored by or
|
|
8
|
+
endorsed by Google. The model weights are not part of this repository; Google distributes them
|
|
9
|
+
under the Apache License, Version 2.0. "Gemma" and "EmbeddingGemma" are trademarks of Google LLC.
|
|
10
|
+
|
|
11
|
+
The reference behaviour it reproduces is that of the Hugging Face transformers and
|
|
12
|
+
sentence-transformers implementations (Apache License, Version 2.0).
|