return_tensors

Author: SeBorgey

src/autointent/_wrappers/embedder.py

@@ -235,15 +235,18 @@ def load(cls, path: Path | str, override_config: EmbedderConfig | None = None) -
 
         return cls(EmbedderConfig(**kwargs))
 
-    def embed(self, utterances: list[str], task_type: TaskTypeEnum | None = None) -> npt.NDArray[np.float32]:
+    def embed(
+        self, utterances: list[str], task_type: TaskTypeEnum | None = None, return_tensors: bool = False
+    ) -> npt.NDArray[np.float32] | torch.Tensor:
         """Calculate embeddings for a list of utterances.
 
         Args:
             utterances: List of input texts to calculate embeddings for.
             task_type: Type of task for which embeddings are calculated.
+            return_tensors: If True, return a PyTorch tensor; otherwise, return a numpy array.
 
         Returns:
-            A numpy array of embeddings.
+            A numpy array or PyTorch tensor of embeddings.
         """
         if len(utterances) == 0:
             msg = "Empty input"
@@ -263,7 +266,10 @@ def embed(self, utterances: list[str], task_type: TaskTypeEnum | None = None) ->
             embeddings_path = _get_embeddings_path(hasher.hexdigest())
             if embeddings_path.exists():
                 logger.debug("loading embeddings from %s", str(embeddings_path))
-                return np.load(embeddings_path)  # type: ignore[no-any-return]
+                embeddings_np = np.load(embeddings_path)
+                if return_tensors:
+                    return torch.from_numpy(embeddings_np).to(self.config.device)
+                return embeddings_np  # type: ignore[no-any-return]
 
         self._model = self._load_model()
 
@@ -281,15 +287,19 @@ def embed(self, utterances: list[str], task_type: TaskTypeEnum | None = None) ->
 
         embeddings = self._model.encode(
             utterances,
-            convert_to_numpy=True,
+            convert_to_numpy=not return_tensors,
+            convert_to_tensor=return_tensors,
             batch_size=self.config.batch_size,
             normalize_embeddings=True,
             prompt=prompt,
         )
 
         if self.config.use_cache:
+            embeddings_to_save = embeddings
+            if return_tensors:
+                embeddings_to_save = embeddings.cpu().numpy()
             embeddings_path.parent.mkdir(parents=True, exist_ok=True)
-            np.save(embeddings_path, embeddings)
+            np.save(embeddings_path, embeddings_to_save)
 
         return embeddings
 

src/autointent/modules/scoring/_gcn/gcn_scorer.py

@@ -132,13 +132,13 @@ def fit(self, utterances: list[str], labels: ListOfLabels, descriptions: list[st
         self._embedder = Embedder(self.embedder_config)
         self._label_embedder = Embedder(self.label_embedder_config)
 
-        x_tensor = torch.tensor(self._embedder.embed(utterances, TaskTypeEnum.classification))
+        x_tensor = self._embedder.embed(utterances, TaskTypeEnum.classification, return_tensors=True)
         y_tensor_dtype = torch.float if self._multilabel else torch.long
         y_tensor = torch.tensor(labels, dtype=y_tensor_dtype)
 
-        label_embeddings = torch.tensor(self._label_embedder.embed(descriptions, TaskTypeEnum.classification)).to(
-            self.torch_config.device
-        )
+        label_embeddings = self._label_embedder.embed(
+            descriptions, TaskTypeEnum.classification, return_tensors=True
+        ).to(self.torch_config.device)
 
         self._model = TextMLGCN(
             num_classes=self._n_classes,
@@ -169,7 +169,7 @@ def predict(self, utterances: list[str]) -> npt.NDArray[Any]:
         if not hasattr(self, "_model"):
             msg = "Model is not trained. Call fit() first."
             raise RuntimeError(msg)
-        x_tensor = torch.tensor(self._embedder.embed(utterances, TaskTypeEnum.classification))
+        x_tensor = self._embedder.embed(utterances, TaskTypeEnum.classification, return_tensors=True)
         return self._predict_tensors(x_tensor)
 
     def clear_cache(self) -> None: