rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-01-04 19:02:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(python): Add gemini text embedding function (#806)

Browse Source 
Named it Gemini-text for now. Not sure how complicated it will be to
support both text and multimodal embeddings under the same class
"gemini"..But its not something to worry about for now I guess.
This commit is contained in:
Ayush Chaurasia
2024-01-13 12:08:55 +05:30
committed by GitHub
parent 986891db98
commit 4568df422d
4 changed files with 187 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
python/tests/test_embeddings_slow.py
View File

@@ -89,7 +89,7 @@ def test_openclip(tmp_path):
db = lancedb.connect(tmp_path)
registry = get_registry()
func = registry.get("open-clip").create()
func = registry.get("open-clip").create(max_retries=0)
class Images(LanceModel):
label: str
@@ -170,7 +170,7 @@ def test_cohere_embedding_function():
@pytest.mark.slow
def test_instructor_embedding(tmp_path):
model = get_registry().get("instructor").create()
model = get_registry().get("instructor").create(max_retries=0)
class TextModel(LanceModel):
text: str = model.SourceField()
@@ -182,3 +182,23 @@ def test_instructor_embedding(tmp_path):
tbl.add(df)
assert len(tbl.to_pandas()["vector"][0]) == model.ndims()
@pytest.mark.slow
@pytest.mark.skipif(
os.environ.get("GOOGLE_API_KEY") is None, reason="GOOGLE_API_KEY not set"
)
def test_gemini_embedding(tmp_path):
model = get_registry().get("gemini-text").create(max_retries=0)
class TextModel(LanceModel):
text: str = model.SourceField()
vector: Vector(model.ndims()) = model.VectorField()
df = pd.DataFrame({"text": ["hello world", "goodbye world"]})
db = lancedb.connect(tmp_path)
tbl = db.create_table("test", schema=TextModel, mode="overwrite")
tbl.add(df)
assert len(tbl.to_pandas()["vector"][0]) == model.ndims()
assert tbl.search("hello").limit(1).to_pandas()["text"][0] == "hello world"