From 1946d00cd5854b03542be6fc4a78743cf1f1389c Mon Sep 17 00:00:00 2001 From: bhakarboy01 Date: Sun, 16 Mar 2025 03:55:23 +0530 Subject: [PATCH] added consistency in embed model names --- google/generativeai/embedding.py | 2 +- samples/embed.py | 4 ++-- tests/test_embedding.py | 2 +- tests/test_embedding_async.py | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/google/generativeai/embedding.py b/google/generativeai/embedding.py index 15645c792..f93f3371a 100644 --- a/google/generativeai/embedding.py +++ b/google/generativeai/embedding.py @@ -28,7 +28,7 @@ from google.generativeai.types import text_types from google.generativeai.types import content_types -DEFAULT_EMB_MODEL = "models/embedding-001" +DEFAULT_EMB_MODEL = "embedding-001" EMBEDDING_MAX_BATCH_SIZE = 100 EmbeddingTaskType = protos.TaskType diff --git a/samples/embed.py b/samples/embed.py index a5897639b..ce5e3cf0c 100644 --- a/samples/embed.py +++ b/samples/embed.py @@ -22,7 +22,7 @@ def test_embed_content(self): text = "Hello World!" result = genai.embed_content( - model="models/text-embedding-004", content=text, output_dimensionality=10 + model="text-embedding-004", content=text, output_dimensionality=10 ) print(result["embedding"]) # [END embed_content] @@ -37,7 +37,7 @@ def batch_embed_contents(self): "How does the brain work?", ] result = genai.embed_content( - model="models/text-embedding-004", content=texts, output_dimensionality=10 + model="text-embedding-004", content=texts, output_dimensionality=10 ) print(result) # [END batch_embed_contents] diff --git a/tests/test_embedding.py b/tests/test_embedding.py index a208a4743..b149890c2 100644 --- a/tests/test_embedding.py +++ b/tests/test_embedding.py @@ -26,7 +26,7 @@ from absl.testing import absltest from absl.testing import parameterized -DEFAULT_EMB_MODEL = "models/embedding-001" +DEFAULT_EMB_MODEL = "embedding-001" class UnitTests(parameterized.TestCase): diff --git a/tests/test_embedding_async.py b/tests/test_embedding_async.py index 367cf7ded..128788729 100644 --- a/tests/test_embedding_async.py +++ b/tests/test_embedding_async.py @@ -26,7 +26,7 @@ from absl.testing import absltest from absl.testing import parameterized -DEFAULT_EMB_MODEL = "models/embedding-001" +DEFAULT_EMB_MODEL = "embedding-001" class AsyncTests(parameterized.TestCase, unittest.IsolatedAsyncioTestCase):