Update optimum_encoder.py

optimum_encoder.py

@@ -82,7 +82,7 @@ class OptimumEncoder(BaseEncoder):
         short_text = ["short"]
         short_encoded_input = tokenizer(
             short_text, padding=True, truncation=True, return_tensors="pt"
-        ).to("cuda")
+        ).to(self.device)
         short_output = ort_model(**short_encoded_input)
         
         print("Building engine for a long sequence...")
@@ -123,11 +123,11 @@ class OptimumEncoder(BaseEncoder):
 
             if pooling_strategy == "mean":
                 embeddings = self._mean_pooling(
-                    model_output, encoded_input["attention_mask"], convert_to_numpy
+                    model_output, encoded_input["attention_mask"]
                 )
             elif pooling_strategy == "max":
                 embeddings = self._max_pooling(
-                    model_output, encoded_input["attention_mask"], convert_to_numpy
+                    model_output, encoded_input["attention_mask"]
                 )
             else:
                 raise ValueError(
@@ -135,36 +135,30 @@ class OptimumEncoder(BaseEncoder):
                 )
 
             if normalize_embeddings:
-                if convert_to_numpy:
-                    embeddings = normalize(embeddings[:, 0]).astype(np.float32)
-                else:
-                    embeddings = self._torch.nn.functional.normalize(embeddings, p=2, dim=1).detach().cpu().tolist()
+                    embeddings = self._torch.nn.functional.normalize(embeddings, p=2, dim=1)
+                
+            if convert_to_numpy:
+                embeddings.detach().cpu().numpy()
+            else:
+                embeddings.tolist()
                 
             all_embeddings.extend(embeddings)
 
         return all_embeddings
 
-    def _mean_pooling(self, model_output, attention_mask, convert_to_numpy):
+    def _mean_pooling(self, model_output, attention_mask):
         token_embeddings = model_output[0]
         input_mask_expanded = (
             attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
         )
-        embeddings = self._torch.sum(
+        return self._torch.sum(
             token_embeddings * input_mask_expanded, 1
         ) / self._torch.clamp(input_mask_expanded.sum(1), min=1e-9)
-        if convert_to_numpy:
-            return embeddings.detach().cpu().numpy()
-        else:
-            return embeddings
 
-    def _max_pooling(self, model_output, attention_mask, convert_to_numpy):
+    def _max_pooling(self, model_output, attention_mask):
         token_embeddings = model_output[0]
         input_mask_expanded = (
             attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
         )
         token_embeddings[input_mask_expanded == 0] = -1e9
-        embeddings = self._torch.max(token_embeddings, 1)[0]
-        if convert_to_numpy:
-            return embeddings.detach().cpu().numpy()
-        else:
-            return embeddings
+        return self._torch.max(token_embeddings, 1)[0]