AI/ragflow
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/infiniflow/ragflow.git synced 2025-08-14 06:05:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix gemini embedding error. (#4356)

Browse Source 
### What problem does this PR solve?

#4314

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)
This commit is contained in:
Kevin Hu 2025-01-06 14:41:29 +08:00 committed by GitHub
parent 983ec0666c
commit b93c136797
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
rag/llm/embedding_model.py
View File

@ -490,6 +490,7 @@ class BedrockEmbed(Base):
return np.array(embeddings), token_count return np.array(embeddings), token_count
class GeminiEmbed(Base): class GeminiEmbed(Base):
def __init__(self, key, model_name='models/text-embedding-004', def __init__(self, key, model_name='models/text-embedding-004',
**kwargs): **kwargs):
@ -505,7 +506,7 @@ class GeminiEmbed(Base):
for i in range(0, len(texts), batch_size): for i in range(0, len(texts), batch_size):
result = genai.embed_content( result = genai.embed_content(
model=self.model_name, model=self.model_name,
content=texts[i, i + batch_size], content=texts[i: i + batch_size],
task_type="retrieval_document", task_type="retrieval_document",
title="Embedding of single string") title="Embedding of single string")
ress.extend(result['embedding']) ress.extend(result['embedding'])
@ -519,7 +520,8 @@ class GeminiEmbed(Base):
task_type="retrieval_document", task_type="retrieval_document",
title="Embedding of single string") title="Embedding of single string")
token_count = num_tokens_from_string(text) token_count = num_tokens_from_string(text)
return np.array(result['embedding']),token_count return np.array(result['embedding']), token_count
class NvidiaEmbed(Base): class NvidiaEmbed(Base):
def __init__( def __init__(