fix: StopIteration issue

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)