mirror of
https://github.com/zenlm/zen-embedding-2.git
synced 2026-10-30 15:06:52 +00:00
zeekay
e5160ea677
One 768-d space for text, code, images, video and audio. The weights are Google DeepMind's EmbeddingGemma 2 at 914f7f89, Apache-2.0, published at huggingface.co/zenlm/zen-embedding-2; NOTICE carries the attribution.
21 lines
1000 B
Python
21 lines
1000 B
Python
"""Text, code and an image in one vector space: rank three documents for a query.
|
|
|
|
uv run --with "sentence-transformers[image]" --with torch examples/encode.py
|
|
"""
|
|
import torch
|
|
from sentence_transformers import SentenceTransformer
|
|
|
|
dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float32
|
|
model = SentenceTransformer("zenlm/zen-embedding-2", model_kwargs={"torch_dtype": dtype},
|
|
config_kwargs={"audio_config": None})
|
|
|
|
query = model.encode("task: code retrieval | query: where is the cache key built?")
|
|
docs = [
|
|
"title: cache.py | text: def make_key(args, kwargs): return hash((args, tuple(sorted(kwargs.items()))))",
|
|
"title: http.py | text: def fetch(url): return requests.get(url).text",
|
|
"title: none | text: Release notes for version 2.4.",
|
|
]
|
|
scores = model.similarity(query, model.encode(docs))[0]
|
|
for score, doc in sorted(zip(scores.tolist(), docs), reverse=True):
|
|
print(f"{score:.3f} {doc[:60]}")
|