refactor(labeler): use set_embeddings in labeler (#165)

jina-ai · Oct 23, 2021 · d8d875f · d8d875f
1 parent 0d8e0b5
commit d8d875f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 28 deletions.
diff --git a/finetuner/labeler/executor.py b/finetuner/labeler/executor.py
@@ -4,7 +4,7 @@
 from jina import Executor, DocumentArray, requests, DocumentArrayMemmap
 from jina.helper import cached_property
 
-from ..helper import get_framework
+from ..embedding import set_embeddings
 from ..tuner import fit, save
 
 
@@ -34,37 +34,15 @@ def embed(self, docs: DocumentArray, parameters: Dict, **kwargs):
         if not docs:
             return
         self._all_data.reload()
-        da = self._all_data.sample(
+        _catalog = self._all_data.sample(
             min(len(self._all_data), int(parameters.get('sample_size', 1000)))
         )
 
-        f_type = get_framework(self._embed_model)
-
-        if f_type == 'keras':
-            da_input = da.blobs
-            docs_input = docs.blobs
-            da.embeddings = self._embed_model(da_input).numpy()
-            docs.embeddings = self._embed_model(docs_input).numpy()
-        elif f_type == 'torch':
-            import torch
-
-            self._embed_model.eval()
-            da_input = torch.from_numpy(da.blobs)
-            docs_input = torch.from_numpy(docs.blobs)
-            with torch.inference_mode():
-                da.embeddings = self._embed_model(da_input).detach().numpy()
-                docs.embeddings = self._embed_model(docs_input).detach().numpy()
-        elif f_type == 'paddle':
-            import paddle
-
-            self._embed_model.eval()
-            da_input = paddle.to_tensor(da.blobs)
-            docs_input = paddle.to_tensor(docs.blobs)
-            da.embeddings = self._embed_model(da_input).detach().numpy()
-            docs.embeddings = self._embed_model(docs_input).detach().numpy()
+        set_embeddings(docs, self._embed_model)
+        set_embeddings(_catalog, self._embed_model)
 
         docs.match(
-            da,
+            _catalog,
             metric=self._metric,
             limit=int(parameters.get('topk', 10)),
             exclude_self=True,

diff --git a/finetuner/toydata.py b/finetuner/toydata.py
@@ -300,12 +300,13 @@ def _download_fashion_doc(
     ):
 
         _d = Document(
-            content=(raw_img / 255.0).astype(np.float32),
+            content=raw_img,
             tags={
                 'class': int(lbl),
             },
         )
         _d.convert_image_blob_to_uri()
+        _d.blob = (_d.blob / 255.0).astype(np.float32)
         yield _d