Ver código fonte

fix SILICONFLOW embedding error (#2363)

### What problem does this PR solve?

#2335  fix SILICONFLOW embedding error

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

---------

Co-authored-by: Zhedong Cen <cenzhedong2@126.com>
tags/v0.11.0
黄腾 1 ano atrás
pai
commit
35b7d17d97
Nenhuma conta vinculada ao e-mail do autor do commit
1 arquivos alterados com 33 adições e 4 exclusões
  1. 33
    4
      rag/llm/embedding_model.py

+ 33
- 4
rag/llm/embedding_model.py Ver arquivo

@@ -577,11 +577,40 @@ class UpstageEmbed(OpenAIEmbed):
super().__init__(key, model_name, base_url)


class SILICONFLOWEmbed(OpenAIEmbed):
def __init__(self, key, model_name, base_url="https://api.siliconflow.cn/v1"):
class SILICONFLOWEmbed(Base):
def __init__(
self, key, model_name, base_url="https://api.siliconflow.cn/v1/embeddings"
):
if not base_url:
base_url = "https://api.siliconflow.cn/v1"
super().__init__(key, model_name, base_url)
base_url = "https://api.siliconflow.cn/v1/embeddings"
self.headers = {
"accept": "application/json",
"content-type": "application/json",
"authorization": f"Bearer {key}",
}
self.base_url = base_url
self.model_name = model_name

def encode(self, texts: list, batch_size=32):
payload = {
"model": self.model_name,
"input": texts,
"encoding_format": "float",
}
res = requests.post(self.base_url, json=payload, headers=self.headers).json()
return (
np.array([d["embedding"] for d in res["data"]]),
res["usage"]["total_tokens"],
)

def encode_queries(self, text):
payload = {
"model": self.model_name,
"input": text,
"encoding_format": "float",
}
res = requests.post(self.base_url, json=payload, headers=self.headers).json()
return np.array(res["data"][0]["embedding"]), res["usage"]["total_tokens"]


class ReplicateEmbed(Base):

Carregando…
Cancelar
Salvar