Skip to content

Commit

Permalink
add batch embed to avoid empty response from server issue
Browse files Browse the repository at this point in the history
  • Loading branch information
gary fan committed Feb 14, 2024
1 parent c87527c commit ab82763
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 ab82763

Please sign in to comment.