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__(
|