From b93c136797d99db8aff4ac3c80993b2ab20b88d1 Mon Sep 17 00:00:00 2001 From: Kevin Hu Date: Mon, 6 Jan 2025 14:41:29 +0800 Subject: [PATCH] Fix gemini embedding error. (#4356) ### What problem does this PR solve? #4314 ### Type of change - [x] Bug Fix (non-breaking change which fixes an issue) --- rag/llm/embedding_model.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/rag/llm/embedding_model.py b/rag/llm/embedding_model.py index 785650a57..182d4cf7a 100644 --- a/rag/llm/embedding_model.py +++ b/rag/llm/embedding_model.py @@ -490,6 +490,7 @@ class BedrockEmbed(Base): return np.array(embeddings), token_count + class GeminiEmbed(Base): def __init__(self, key, model_name='models/text-embedding-004', **kwargs): @@ -505,7 +506,7 @@ class GeminiEmbed(Base): for i in range(0, len(texts), batch_size): result = genai.embed_content( model=self.model_name, - content=texts[i, i + batch_size], + content=texts[i: i + batch_size], task_type="retrieval_document", title="Embedding of single string") ress.extend(result['embedding']) @@ -519,7 +520,8 @@ class GeminiEmbed(Base): task_type="retrieval_document", title="Embedding of single string") token_count = num_tokens_from_string(text) - return np.array(result['embedding']),token_count + return np.array(result['embedding']), token_count + class NvidiaEmbed(Base): def __init__(