Draft demo

client/clip_client/client.py

@@ -11,6 +11,7 @@
     Generator,
     Iterable,
     Dict,
+    List,
 )
 from urllib.parse import urlparse
 from functools import partial
@@ -61,6 +62,7 @@ def __init__(self, server: str, credential: dict = {}, **kwargs):
             raise ValueError(f'{server} is not a valid scheme')
 
         self._authorization = credential.get('Authorization', None)
+        self._batch_size = 128
 
     @overload
     def encode(
@@ -114,9 +116,16 @@ def encode(self, content, **kwargs):
             total=len(content) if hasattr(content, '__len__') else None,
         )
         results = DocumentArray()
+
+        parameters = kwargs.get('parameters', {})
+        model_name = parameters.get('model', '')
+        payload = self._get_post_parameters(content, kwargs)
+        payload.update(on=f'/encode/{model_name}'.rstrip('/'))
+        payload.update(inputs=self._iter_doc(content))
+
         with self._pbar:
             self._client.post(
-                **self._get_post_payload(content, kwargs),
+                **payload,
                 on_done=partial(self._gather_result, results=results),
             )
         return self._unboxed_result(results)
@@ -188,12 +197,8 @@ def _iter_doc(self, content) -> Generator['Document', None, None]:
                     ),
                 )
 
-    def _get_post_payload(self, content, kwargs):
-        parameters = kwargs.get('parameters', {})
-        model_name = parameters.get('model', '')
+    def _get_post_parameters(self, content, kwargs):
         payload = dict(
-            on=f'/encode/{model_name}'.rstrip('/'),
-            inputs=self._iter_doc(content),
             request_size=kwargs.get('batch_size', 8),
             total_docs=len(content) if hasattr(content, '__len__') else None,
         )
@@ -284,9 +289,14 @@ async def aencode(self, content, **kwargs):
         )
 
         results = DocumentArray()
-        async for da in self._async_client.post(
-            **self._get_post_payload(content, kwargs)
-        ):
+
+        parameters = kwargs.get('parameters', {})
+        model_name = parameters.get('model', '')
+        payload = self._get_post_parameters(content, kwargs)
+        payload.update(on=f'/encode/{model_name}'.rstrip('/'))
+        payload.update(inputs=self._iter_doc(content))
+
+        async for da in self._async_client.post(**payload):
             if not results:
                 self._pbar.start_task(self._r_task)
             results.extend(da)
@@ -369,23 +379,6 @@ def _iter_rank_docs(
                     ),
                 )
 
-    def _get_rank_payload(self, content, kwargs):
-        parameters = kwargs.get('parameters', {})
-        model_name = parameters.get('model', '')
-        payload = dict(
-            on=f'/rank/{model_name}'.rstrip('/'),
-            inputs=self._iter_rank_docs(
-                content, _source=kwargs.get('source', 'matches')
-            ),
-            request_size=kwargs.get('batch_size', 8),
-            total_docs=len(content) if hasattr(content, '__len__') else None,
-        )
-        if self._scheme == 'grpc' and self._authorization:
-            payload.update(metadata=('authorization', self._authorization))
-        elif self._scheme == 'http' and self._authorization:
-            payload.update(headers={'Authorization': self._authorization})
-        return payload
-
     def rank(self, docs: Iterable['Document'], **kwargs) -> 'DocumentArray':
         """Rank image-text matches according to the server CLIP model.
         Given a Document with nested matches, where the root is image/text and the matches is in another modality, i.e.
@@ -400,9 +393,18 @@ def rank(self, docs: Iterable['Document'], **kwargs) -> 'DocumentArray':
             total=len(docs),
         )
         results = DocumentArray()
+
+        parameters = kwargs.get('parameters', {})
+        model_name = parameters.get('model', '')
+        payload = self._get_post_parameters(docs, kwargs)
+        payload.update(on=f'/rank/{model_name}'.rstrip('/'))
+        payload.update(
+            inputs=self._iter_rank_docs(docs, _source=kwargs.get('source', 'matches'))
+        )
+
         with self._pbar:
             self._client.post(
-                **self._get_rank_payload(docs, kwargs),
+                **payload,
                 on_done=partial(self._gather_result, results=results),
             )
         return results
@@ -415,7 +417,16 @@ async def arank(self, docs: Iterable['Document'], **kwargs) -> 'DocumentArray':
             total=len(docs),
         )
         results = DocumentArray()
-        async for da in self._async_client.post(**self._get_rank_payload(docs, kwargs)):
+
+        parameters = kwargs.get('parameters', {})
+        model_name = parameters.get('model', '')
+        payload = self._get_post_parameters(docs, kwargs)
+        payload.update(on=f'/rank/{model_name}'.rstrip('/'))
+        payload.update(
+            inputs=self._iter_rank_docs(docs, _source=kwargs.get('source', 'matches'))
+        )
+
+        async for da in self._async_client.post(**payload):
             if not results:
                 self._pbar.start_task(self._r_task)
             results.extend(da)
@@ -428,3 +439,54 @@ async def arank(self, docs: Iterable['Document'], **kwargs) -> 'DocumentArray':
             )
 
         return results
+
+    def index(self, content: Iterable['Document'], **kwargs):
+        """Index the embeddings created by server CLIP model.
+        Given the document with embeddings, this function create an indexer which index
+        the embeddings. This will be used for top k search. ``AnnLiteIndexer`` is used
+        by default.
+        :param content: docs to be indexed.
+        """
+        self._prepare_streaming(
+            not kwargs.get('show_progress'),
+            total=len(content) if hasattr(content, '__len__') else None,
+        )
+
+        payload = self._get_post_parameters(content, kwargs)
+        payload.update(on='/index')
+        payload.update(inputs=self._iter_doc(content))
+
+        with self._pbar:
+            self._client.post(
+                **payload,
+            )
+
+    def search(self, content: List[str], **kwargs) -> DocumentArray:
+        """Search for top k results for given query string or ``Document``.
+        If the input is a string, will use this string as query. If the input is a
+        ``Document``, will use the ``text`` field as query.
+        :param content: list of queries.
+        :return: top limit results.
+        """
+        self._prepare_streaming(
+            not kwargs.get('show_progress'),
+            total=len(content) if hasattr(content, '__len__') else None,
+        )
+        results = DocumentArray()
+
+        payload = self._get_post_parameters(content, kwargs)
+        payload.update(on='/search')
+        payload.update(inputs=self._iter_doc(content))
+
+        with self._pbar:
+            self._client.post(
+                **payload,
+                on_done=partial(self._gather_result, results=results),
+            )
+        return results
+
+    def status(self):
+        payload = dict(
+            on='/status',
+        )
+        return self._client.post(**payload)

server/clip_server/demo-flow.yml

@@ -0,0 +1,21 @@
+jtype: Flow
+version: '1'
+with:
+  port: 61000
+executors:
+  - name: encoder
+    uses:
+      jtype: CLIPEncoder
+      metas:
+        py_modules:
+          - clip_server.executors.clip_torch
+    replicas: 1
+  - name: indexer
+    uses:
+      jtype: AnnLiteIndexer
+      with:
+        dim: 512
+      metas:
+        py_modules:
+          - annlite.executor
+    replicas: 1