api inference compat response

Signed-off-by: Raphael Glon <[email protected]>

Author: oOraph

src/huggingface_inference_toolkit/env_utils.py

@@ -1,3 +1,6 @@
+import os
+
+
 def strtobool(val: str) -> bool:
     """Convert a string representation of truth to True or False booleans.
     True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values
@@ -20,3 +23,7 @@ def strtobool(val: str) -> bool:
     raise ValueError(
         f"Invalid truth value, it should be a string but {val} was provided instead."
     )
+
+
+def api_inference_compat():
+    return strtobool(os.getenv("API_INFERENCE_COMPAT", "false"))

src/huggingface_inference_toolkit/sentence_transformers_utils.py

@@ -1,6 +1,9 @@
 import importlib.util
+import os
 from typing import Any, Dict, List, Tuple, Union
 
+from huggingface_inference_toolkit.env_utils import api_inference_compat
+
 try:
     from typing import Literal
 except ImportError:
@@ -26,7 +29,10 @@ def __call__(self, source_sentence: str, sentences: List[str]) -> Dict[str, floa
         embeddings1 = self.model.encode(source_sentence, convert_to_tensor=True)
         embeddings2 = self.model.encode(sentences, convert_to_tensor=True)
         similarities = util.pytorch_cos_sim(embeddings1, embeddings2).tolist()[0]
-        return {"similarities": similarities}
+        if api_inference_compat():
+            return similarities
+        else:
+            return {"similarities": similarities}
 
 
 class SentenceEmbeddingPipeline:
@@ -36,7 +42,10 @@ def __init__(self, model_dir: str, device: Union[str, None] = None, **kwargs: An
 
     def __call__(self, sentences: Union[str, List[str]]) -> Dict[str, List[float]]:
         embeddings = self.model.encode(sentences).tolist()
-        return {"embeddings": embeddings}
+        if api_inference_compat():
+            return embeddings
+        else:
+            return {"embeddings": embeddings}
 
 
 class SentenceRankingPipeline:

src/huggingface_inference_toolkit/utils.py

@@ -1,4 +1,5 @@
 import importlib.util
+import os
 import sys
 from pathlib import Path
 from typing import Optional, Union