Merge branch 'dev' into monitoring

Author: andhreljaKern

.drone.yml

@@ -67,7 +67,7 @@ name: arm64
 platform:
   arch: arm64
 
-steps:  
+steps:
   - name: submodules
     image: alpine/git
     commands:
@@ -96,41 +96,41 @@ trigger:
     - push
 
 ---
-kind: pipeline
-type: docker
-name: gpu
-
-platform:
-  arch: amd64
-
-steps:
-  - name: submodules
-    image: alpine/git
-    commands:
-      - git submodule init
-      - 'git config --global url."https://github.com/".insteadOf [email protected]:'
-      - "git submodule update --recursive"
-  - name: build and publish
-    image: plugins/docker
-    settings:
-      registry: registry.dev.kern.ai
-      username:
-        from_secret: docker_username
-      password:
-        from_secret: docker_password
-      dockerfile: gpu.Dockerfile
-      repo: "registry.dev.kern.ai/${DRONE_REPO}"
-      tags: ["${DRONE_COMMIT_SHA}-gpu", "${DRONE_COMMIT_BRANCH}-gpu"]
-      cache_from:
-          - "registry.dev.kern.ai/${DRONE_REPO}:${DRONE_COMMIT_BRANCH}-gpu"
-
-trigger:
-  branch:
-    - dev
-  event:
-    - push
-
----
+# kind: pipeline
+# type: docker
+# name: gpu
+
+# platform:
+#   arch: amd64
+
+# steps:
+#   - name: submodules
+#     image: alpine/git
+#     commands:
+#       - git submodule init
+#       - 'git config --global url."https://github.com/".insteadOf [email protected]:'
+#       - "git submodule update --recursive"
+#   - name: build and publish
+#     image: plugins/docker
+#     settings:
+#       registry: registry.dev.kern.ai
+#       username:
+#         from_secret: docker_username
+#       password:
+#         from_secret: docker_password
+#       dockerfile: gpu.Dockerfile
+#       repo: "registry.dev.kern.ai/${DRONE_REPO}"
+#       tags: ["${DRONE_COMMIT_SHA}-gpu", "${DRONE_COMMIT_BRANCH}-gpu"]
+#       cache_from:
+#           - "registry.dev.kern.ai/${DRONE_REPO}:${DRONE_COMMIT_BRANCH}-gpu"
+
+# trigger:
+#   branch:
+#     - dev
+#   event:
+#     - push
+
+# ---
 kind: pipeline
 type: docker
 name: amd64-dockerhub
@@ -160,12 +160,42 @@ trigger:
     - tag
 
 ---
+# kind: pipeline
+# type: docker
+# name: amd64-gpu-dockerhub
+
+# platform:
+#   arch: amd64
+
+# steps:
+#   - name: submodules
+#     image: alpine/git
+#     commands:
+#       - git submodule init
+#       - 'git config --global url."https://github.com/".insteadOf [email protected]:'
+#       - "git submodule update --recursive"
+#   - name: build and publish
+#     image: plugins/docker
+#     settings:
+#       username:
+#         from_secret: dockerhub_username
+#       password:
+#         from_secret: dockerhub_password
+#       dockerfile: gpu.Dockerfile
+#       repo: "kernai/${DRONE_REPO_NAME}"
+#       tag: "${DRONE_TAG}-gpu"
+
+# trigger:
+#   event:
+#     - tag
+
+# ---
 kind: pipeline
 type: docker
-name: amd64-gpu-dockerhub
+name: arm64-dockerhub
 
 platform:
-  arch: amd64
+  arch: arm64
 
 steps:
   - name: submodules
@@ -181,57 +211,27 @@ steps:
         from_secret: dockerhub_username
       password:
         from_secret: dockerhub_password
-      dockerfile: gpu.Dockerfile
       repo: "kernai/${DRONE_REPO_NAME}"
-      tag: "${DRONE_TAG}-gpu"
+      tag: "${DRONE_TAG}-drone-arm64"
 
 trigger:
   event:
     - tag
 
 ---
 kind: pipeline
-type: docker
-name: arm64-dockerhub
-
-platform:
-  arch: arm64
-
+name: manifest-version
 steps:
-  - name: submodules
-    image: alpine/git
-    commands:
-      - git submodule init
-      - 'git config --global url."https://github.com/".insteadOf [email protected]:'
-      - "git submodule update --recursive"
-  - name: build and publish
-    image: plugins/docker
+  - name: manifest
+    image: plugins/manifest
     settings:
+      spec: drone-manifest-version.tmpl
+      tag: "${DRONE_TAG}"
+      ignore_missing: true
       username:
         from_secret: dockerhub_username
       password:
         from_secret: dockerhub_password
-      repo: "kernai/${DRONE_REPO_NAME}"
-      tag: "${DRONE_TAG}-drone-arm64"
-
-trigger:
-  event:
-    - tag
-
----
-kind: pipeline
-name: manifest-version
-steps:
-- name: manifest
-  image: plugins/manifest
-  settings:
-    spec: drone-manifest-version.tmpl
-    tag: "${DRONE_TAG}"
-    ignore_missing: true
-    username:
-      from_secret: dockerhub_username
-    password:
-      from_secret: dockerhub_password
 
 depends_on:
   - amd64-dockerhub
@@ -245,16 +245,16 @@ trigger:
 kind: pipeline
 name: manifest-latest
 steps:
-- name: manifest
-  image: plugins/manifest
-  settings:
-    spec: drone-manifest-latest.tmpl
-    tag: "${DRONE_TAG}"
-    ignore_missing: true
-    username:
-      from_secret: dockerhub_username
-    password:
-      from_secret: dockerhub_password
+  - name: manifest
+    image: plugins/manifest
+    settings:
+      spec: drone-manifest-latest.tmpl
+      tag: "${DRONE_TAG}"
+      ignore_missing: true
+      username:
+        from_secret: dockerhub_username
+      password:
+        from_secret: dockerhub_password
 
 depends_on:
   - manifest-version

src/embedders/__init__.py

@@ -165,6 +165,7 @@ def load_pca_weights(self, file_name: str):
     def _reduce(
         self,
         documents: List[Union[str, Doc]],
+        as_generator: bool,
         fit_model: bool,
         fit_after_n_batches: int,
     ):
@@ -178,11 +179,11 @@ def _reduce_batch(
         fit_after_n_batches: int,
     ) -> Union[List, Generator]:
         if as_generator:
-            return self._reduce(documents, fit_model, fit_after_n_batches)
+            return self._reduce(documents, as_generator, fit_model, fit_after_n_batches)
         else:
             embeddings = []
             for embedding_batch in self._reduce(
-                documents, fit_model, fit_after_n_batches
+                documents, as_generator, fit_model, fit_after_n_batches
             ):
                 embeddings.extend(embedding_batch)
             return embeddings

src/embedders/classification/reduce.py

@@ -21,6 +21,7 @@ def _transform(
     def _reduce(
         self,
         documents: List[Union[str, Doc]],
+        as_generator: bool,
         fit_model: bool,
         fit_after_n_batches: int,
     ) -> Generator[List[List[Union[float, int]]], None, None]:
@@ -56,8 +57,16 @@ def _reduce(
                 if batch_idx > fit_after_n_batches:
                     yield self._transform(batch)
         else:
-            embeddings = self.embedder.transform(documents)
-            yield self._transform(embeddings)
+            if as_generator:
+                embeddings = [
+                    emb
+                    for batch in self.embedder.transform(documents, as_generator)
+                    for emb in batch
+                ]
+                yield from util.batch(self._transform(embeddings), self.batch_size)
+            else:
+                embeddings = self.embedder.transform(documents)
+                yield self._transform(embeddings)
 
     @staticmethod
     def load(embedder: dict) -> "PCASentenceReducer":

src/embedders/extraction/reduce.py

@@ -1,3 +1,4 @@
+from spacy.tokens.doc import Doc
 from typing import List, Generator, Union
 import numpy as np
 from src.embedders import PCAReducer, util
@@ -24,7 +25,11 @@ def _transform(
         return batch_unsqueezed
 
     def _reduce(
-        self, documents, fit_model, fit_after_n_batches
+        self,
+        documents: List[Union[str, Doc]],
+        as_generator: bool,
+        fit_model: bool,
+        fit_after_n_batches: int,
     ) -> Generator[List[List[List[Union[float, int]]]], None, None]:
         if fit_model:
             embeddings_training = []
@@ -60,5 +65,13 @@ def _reduce(
                 if batch_idx > fit_after_n_batches:
                     yield self._transform(batch)
         else:
-            embeddings = self.embedder.transform(documents)
-            yield self._transform(embeddings)
+            if as_generator:
+                embeddings = [
+                    emb
+                    for batch in self.embedder.transform(documents, as_generator)
+                    for emb in batch
+                ]
+                yield from util.batch(self._transform(embeddings), self.batch_size)
+            else:
+                embeddings = self.embedder.transform(documents)
+                yield self._transform(embeddings)