Add to feature-extraction + update types

src/huggingface_hub/inference/_client.py

@@ -940,6 +940,7 @@ def feature_extraction(
         text: str,
         *,
         normalize: Optional[bool] = None,
+        prompt_name: Optional[str] = None,
         truncate: Optional[bool] = None,
         model: Optional[str] = None,
     ) -> "np.ndarray":
@@ -956,6 +957,12 @@ def feature_extraction(
             normalize (`bool`, *optional*):
                 Whether to normalize the embeddings or not. Defaults to None.
                 Only available on server powered by Text-Embedding-Inference.
+            prompt_name (`str`, *optional*):
+                The name of the prompt that should be used by for encoding. If not set, no prompt will be applied.
+                Must be a key in the `Sentence Transformers` configuration `prompts` dictionary.
+                For example if ``prompt_name`` is "query" and the ``prompts`` is {"query": "query: ",...},
+                then the sentence "What is the capital of France?" will be encoded as "query: What is the capital of France?"
+                because the prompt text will be prepended before any text to encode.
             truncate (`bool`, *optional*):
                 Whether to truncate the embeddings or not. Defaults to None.
                 Only available on server powered by Text-Embedding-Inference.
@@ -983,6 +990,8 @@ def feature_extraction(
         payload: Dict = {"inputs": text}
         if normalize is not None:
             payload["normalize"] = normalize
+        if prompt_name is not None:
+            payload["prompt_name"] = prompt_name
         if truncate is not None:
             payload["truncate"] = truncate
         response = self.post(json=payload, model=model, task="feature-extraction")

src/huggingface_hub/inference/_generated/_async_client.py

@@ -944,6 +944,7 @@ async def feature_extraction(
         text: str,
         *,
         normalize: Optional[bool] = None,
+        prompt_name: Optional[str] = None,
         truncate: Optional[bool] = None,
         model: Optional[str] = None,
     ) -> "np.ndarray":
@@ -960,6 +961,12 @@ async def feature_extraction(
             normalize (`bool`, *optional*):
                 Whether to normalize the embeddings or not. Defaults to None.
                 Only available on server powered by Text-Embedding-Inference.
+            prompt_name (`str`, *optional*):
+                The name of the prompt that should be used by for encoding. If not set, no prompt will be applied.
+                Must be a key in the `Sentence Transformers` configuration `prompts` dictionary.
+                For example if ``prompt_name`` is "query" and the ``prompts`` is {"query": "query: ",...},
+                then the sentence "What is the capital of France?" will be encoded as "query: What is the capital of France?"
+                because the prompt text will be prepended before any text to encode.
             truncate (`bool`, *optional*):
                 Whether to truncate the embeddings or not. Defaults to None.
                 Only available on server powered by Text-Embedding-Inference.
@@ -988,6 +995,8 @@ async def feature_extraction(
         payload: Dict = {"inputs": text}
         if normalize is not None:
             payload["normalize"] = normalize
+        if prompt_name is not None:
+            payload["prompt_name"] = prompt_name
         if truncate is not None:
             payload["truncate"] = truncate
         response = await self.post(json=payload, model=model, task="feature-extraction")

src/huggingface_hub/inference/_generated/types/feature_extraction.py

@@ -4,11 +4,14 @@
 #   - script: https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-codegen.ts
 #   - specs:  https://github.com/huggingface/huggingface.js/tree/main/packages/tasks/src/tasks.
 from dataclasses import dataclass
-from typing import List, Optional, Union
+from typing import Literal, Optional
 
 from .base import BaseInferenceType
 
 
+FeatureExtractionInputTruncationDirection = Literal["Left", "Right"]
+
+
 @dataclass
 class FeatureExtractionInput(BaseInferenceType):
     """Feature Extraction Input.
@@ -17,6 +20,18 @@ class FeatureExtractionInput(BaseInferenceType):
     https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tei-import.ts.
     """
 
-    inputs: Union[List[str], str]
+    inputs: str
+    """The text to embed."""
     normalize: Optional[bool] = None
+    prompt_name: Optional[str] = None
+    """The name of the prompt that should be used by for encoding. If not set, no prompt
+    will be applied.
+    Must be a key in the `Sentence Transformers` configuration `prompts` dictionary.
+    For example if ``prompt_name`` is "query" and the ``prompts`` is {"query": "query: ",
+    ...},
+    then the sentence "What is the capital of France?" will be encoded as
+    "query: What is the capital of France?" because the prompt text will be prepended before
+    any text to encode.
+    """
     truncate: Optional[bool] = None
+    truncation_direction: Optional["FeatureExtractionInputTruncationDirection"] = None