pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

refine OpenAi Api (#159)

Browse Source
This commit is contained in:
KevinHuSh
2024-03-27 17:55:45 +08:00
committed by GitHub
parent 0cb95c688e
commit bf2e3d7fc1
5 changed files with 16 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
rag/llm/embedding_model.py
View File

@ -139,12 +139,16 @@ class ZhipuEmbed(Base):
self.model_name = model_name
def encode(self, texts: list, batch_size=32):
res = self.client.embeddings.create(input=texts,
arr = []
tks_num = 0
for txt in texts:
res = self.client.embeddings.create(input=txt,
model=self.model_name)
return np.array([d.embedding for d in res.data]
), res.usage.total_tokens
arr.append(res.data[0].embedding)
tks_num += res.usage.total_tokens
return np.array(arr), tks_num
def encode_queries(self, text):
res = self.client.embeddings.create(input=text,
model=self.model_name)
return np.array(res["data"][0]["embedding"]), res.usage.total_tokens
return np.array(res.data[0].embedding), res.usage.total_tokens