diff --git a/client/clip_client/client.py b/client/clip_client/client.py index 5176227ab..262246f82 100644 --- a/client/clip_client/client.py +++ b/client/clip_client/client.py @@ -188,7 +188,6 @@ def _iter_doc( ), ) - d.tags['__ordered_by_CAS__'] = i results.append(d) yield d @@ -212,10 +211,7 @@ def _gather_encode_result(self, response, results: 'DocumentArray'): from rich import filesize r = response.data.docs - for d in r: - index = int(d.tags['__ordered_by_CAS__']) - results[index].embedding = d.embedding - results[index].tags.pop('__ordered_by_CAS__') + results[r[:, 'id']].embeddings = r.embeddings if not self._pbar._tasks[self._r_task].started: self._pbar.start_task(self._r_task) @@ -345,10 +341,7 @@ async def aencode(self, content, **kwargs): async for da in self._async_client.post( **self._get_post_payload(content, results, kwargs) ): - for d in da: - index = int(d.tags['__ordered_by_CAS__']) - results[index].embedding = d.embedding - results[index].tags.pop('__ordered_by_CAS__') + results[da[:, 'id']].embeddings = da.embeddings if not self._pbar._tasks[self._r_task].started: self._pbar.start_task(self._r_task) @@ -395,7 +388,6 @@ def _iter_rank_docs( ), ) - d.tags['__ordered_by_CAS__'] = i results.append(d) yield d @@ -420,10 +412,7 @@ def _gather_rank_result(self, response, results: 'DocumentArray'): from rich import filesize r = response.data.docs - for d in r: - index = int(d.tags['__ordered_by_CAS__']) - results[index].matches = d.matches - results[index].tags.pop('__ordered_by_CAS__') + results[r[:, 'id']][:, 'matches'] = r[:, 'matches'] if not self._pbar._tasks[self._r_task].started: self._pbar.start_task(self._r_task) @@ -513,10 +502,7 @@ async def arank(self, docs: Iterable['Document'], **kwargs) -> 'DocumentArray': async for da in self._async_client.post( **self._get_rank_payload(docs, results, kwargs) ): - for d in da: - index = int(d.tags['__ordered_by_CAS__']) - results[index].matches = d.matches - results[index].tags.pop('__ordered_by_CAS__') + results[da[:, 'id']][:, 'matches'] = da[:, 'matches'] if not self._pbar._tasks[self._r_task].started: self._pbar.start_task(self._r_task)