瀏覽代碼

Refactor: Improve the try catch logic for XinferenceEmbed (#9128)

### What problem does this PR solve?

Improve the try catch logic for XinferenceEmbed

### Type of change


- [x] Refactoring
tags/v0.20.0
Stephen Hu 3 月之前
父節點
當前提交
20b4d88098
沒有連結到貢獻者的電子郵件帳戶。
共有 1 個檔案被更改,包括 4 行新增2 行删除
  1. 4
    2
      rag/llm/embedding_model.py

+ 4
- 2
rag/llm/embedding_model.py 查看文件

@@ -378,8 +378,9 @@ class XinferenceEmbed(Base):
ress = []
total_tokens = 0
for i in range(0, len(texts), batch_size):
res = self.client.embeddings.create(input=texts[i : i + batch_size], model=self.model_name)
res = None
try:
res = self.client.embeddings.create(input=texts[i : i + batch_size], model=self.model_name)
ress.extend([d.embedding for d in res.data])
total_tokens += self.total_token_count(res)
except Exception as _e:
@@ -387,8 +388,9 @@ class XinferenceEmbed(Base):
return np.array(ress), total_tokens

def encode_queries(self, text):
res = self.client.embeddings.create(input=[text], model=self.model_name)
res = None
try:
res = self.client.embeddings.create(input=[text], model=self.model_name)
return np.array(res.data[0].embedding), self.total_token_count(res)
except Exception as _e:
log_exception(_e, res)

Loading…
取消
儲存