Skip to content

Commit

Permalink
feat: add async rerank
Browse files Browse the repository at this point in the history
  • Loading branch information
hanxiao committed Apr 30, 2022
1 parent 20e66b9 commit 12d33c4
Showing 1 changed file with 21 additions and 0 deletions.
21 changes: 21 additions & 0 deletions client/clip_client/client.py
Expand Up @@ -397,3 +397,24 @@ def rerank(self, docs: Iterable['Document'], **kwargs) -> 'DocumentArray':
**self._get_rank_payload(docs, kwargs), on_done=self._gather_result
)
return self._results

async def arerank(self, docs: Iterable['Document'], **kwargs) -> 'DocumentArray':

self._prepare_streaming(
not kwargs.get('show_progress'),
total=len(docs),
)

async for da in self._async_client.post(**self._get_rank_payload(docs, kwargs)):
if not self._results:
self._pbar.start_task(self._r_task)
self._results.extend(da)
self._pbar.update(
self._r_task,
advance=len(da),
total_size=str(
filesize.decimal(int(os.environ.get('JINA_GRPC_RECV_BYTES', '0')))
),
)

return self._results

0 comments on commit 12d33c4

Please sign in to comment.