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)

Files changed (1) hide show
  1. rag/llm/embedding_model.py +4 -2
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, i + batch_size],
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__(