Skip to content

Commit

Permalink
batch提取embedding,避免文本chunk数量过多(>1000)之后报empty response from server 错误 (
Browse files Browse the repository at this point in the history
…#338)

chunk_size较小时,一个文档/知识库一次embedding
request会包括超过千条文本chunks,这时很容易embedding失败。
  • Loading branch information
yaojin3616 committed Feb 23, 2024
2 parents 9ac9442 + ab82763 commit 660f79f
Showing 1 changed file with 24 additions and 9 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -101,16 +101,31 @@ def embed(self, texts: List[str], **kwargs) -> List[List[float]]:
if self.verbose:
print('payload', inp)

max_text_to_split = 200
outp = None
try:
outp = self.client(url=self.url_ep, json=inp, timeout=self.request_timeout).json()
except requests.exceptions.Timeout:
raise Exception(f'timeout in host embedding infer, url=[{self.url_ep}]')
except Exception as e:
raise Exception(f'exception in host embedding infer: [{e}]')

if outp['status_code'] != 200:
raise ValueError(f"API returned an error: {outp['status_message']}")

start_index = 0
len_text = len(texts)
while start_index < len_text:
inp_local = {
'texts':texts[start_index:min(start_index + max_text_to_split, len_text)],
'model':self.model,
'type':emb_type
}
try:
outp_single = self.client(url=self.url_ep, json=inp_local, timeout=self.request_timeout).json()
if outp is None:
outp = outp_single
else:
outp['embeddings'] += outp_single['embeddings']
except requests.exceptions.Timeout:
raise Exception(f'timeout in host embedding infer, url=[{self.url_ep}]')
except Exception as e:
raise Exception(f'exception in host embedding infer: [{e}]')

if outp_single['status_code'] != 200:
raise ValueError(f"API returned an error: {outp['status_message']}")
start_index += max_text_to_split
return outp['embeddings']

def embed_documents(self,
Expand Down

0 comments on commit 660f79f

Please sign in to comment.