Kevin Hu
commited on
Commit
·
6ec0dc3
1
Parent(s):
2820402
Fix gemini embedding error. (#4356)
Browse files### 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
CHANGED
@@ -490,6 +490,7 @@ class BedrockEmbed(Base):
|
|
490 |
|
491 |
return np.array(embeddings), token_count
|
492 |
|
|
|
493 |
class GeminiEmbed(Base):
|
494 |
def __init__(self, key, model_name='models/text-embedding-004',
|
495 |
**kwargs):
|
@@ -505,7 +506,7 @@ class GeminiEmbed(Base):
|
|
505 |
for i in range(0, len(texts), batch_size):
|
506 |
result = genai.embed_content(
|
507 |
model=self.model_name,
|
508 |
-
content=texts[i
|
509 |
task_type="retrieval_document",
|
510 |
title="Embedding of single string")
|
511 |
ress.extend(result['embedding'])
|
@@ -519,7 +520,8 @@ class GeminiEmbed(Base):
|
|
519 |
task_type="retrieval_document",
|
520 |
title="Embedding of single string")
|
521 |
token_count = num_tokens_from_string(text)
|
522 |
-
return np.array(result['embedding']),token_count
|
|
|
523 |
|
524 |
class NvidiaEmbed(Base):
|
525 |
def __init__(
|
|
|
490 |
|
491 |
return np.array(embeddings), token_count
|
492 |
|
493 |
+
|
494 |
class GeminiEmbed(Base):
|
495 |
def __init__(self, key, model_name='models/text-embedding-004',
|
496 |
**kwargs):
|
|
|
506 |
for i in range(0, len(texts), batch_size):
|
507 |
result = genai.embed_content(
|
508 |
model=self.model_name,
|
509 |
+
content=texts[i: i + batch_size],
|
510 |
task_type="retrieval_document",
|
511 |
title="Embedding of single string")
|
512 |
ress.extend(result['embedding'])
|
|
|
520 |
task_type="retrieval_document",
|
521 |
title="Embedding of single string")
|
522 |
token_count = num_tokens_from_string(text)
|
523 |
+
return np.array(result['embedding']), token_count
|
524 |
+
|
525 |
|
526 |
class NvidiaEmbed(Base):
|
527 |
def __init__(
|