From 9a4f84e7b787e84c074cf441a627012f78f8cc61 Mon Sep 17 00:00:00 2001 From: Lucain Pouget Date: Fri, 3 May 2024 16:04:37 +0200 Subject: [PATCH 1/2] Support truncate and normalize in InferenceClient --- src/huggingface_hub/inference/_client.py | 22 +++++++++++++++++-- .../inference/_generated/_async_client.py | 22 +++++++++++++++++-- .../_generated/types/feature_extraction.py | 15 ++++++++----- 3 files changed, 49 insertions(+), 10 deletions(-) diff --git a/src/huggingface_hub/inference/_client.py b/src/huggingface_hub/inference/_client.py index 0c0f079f3a..7122e1149a 100644 --- a/src/huggingface_hub/inference/_client.py +++ b/src/huggingface_hub/inference/_client.py @@ -913,7 +913,14 @@ def document_question_answering( response = self.post(json=payload, model=model, task="document-question-answering") return DocumentQuestionAnsweringOutputElement.parse_obj_as_list(response) - def feature_extraction(self, text: str, *, model: Optional[str] = None) -> "np.ndarray": + def feature_extraction( + self, + text: str, + *, + normalize: Optional[bool] = None, + truncate: Optional[bool] = None, + model: Optional[str] = None, + ) -> "np.ndarray": """ Generate embeddings for a given text. @@ -924,6 +931,12 @@ def feature_extraction(self, text: str, *, model: Optional[str] = None) -> "np.n The model to use for the conversational task. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed Inference Endpoint. If not provided, the default recommended conversational model will be used. Defaults to None. + normalize (`bool`, *optional*): + Whether to normalize the embeddings or not. Defaults to None. + Only available on server powered by Text-Embedding-Inference. + truncate (`bool`, *optional*): + Whether to truncate the embeddings or not. Defaults to None. + Only available on server powered by Text-Embedding-Inference. Returns: `np.ndarray`: The embedding representing the input text as a float32 numpy array. @@ -945,7 +958,12 @@ def feature_extraction(self, text: str, *, model: Optional[str] = None) -> "np.n [ 0.28552425, -0.928395 , -1.2077185 , ..., 0.76810825, -2.1069427 , 0.6236161 ]], dtype=float32) ``` """ - response = self.post(json={"inputs": text}, model=model, task="feature-extraction") + payload = {"inputs": text} + if normalize is not None: + payload["normalize"] = normalize + if truncate is not None: + payload["truncate"] = truncate + response = self.post(json=payload, model=model, task="feature-extraction") np = _import_numpy() return np.array(_bytes_to_dict(response), dtype="float32") diff --git a/src/huggingface_hub/inference/_generated/_async_client.py b/src/huggingface_hub/inference/_generated/_async_client.py index 3cfbdfa1b4..ca8c951319 100644 --- a/src/huggingface_hub/inference/_generated/_async_client.py +++ b/src/huggingface_hub/inference/_generated/_async_client.py @@ -915,7 +915,14 @@ async def document_question_answering( response = await self.post(json=payload, model=model, task="document-question-answering") return DocumentQuestionAnsweringOutputElement.parse_obj_as_list(response) - async def feature_extraction(self, text: str, *, model: Optional[str] = None) -> "np.ndarray": + async def feature_extraction( + self, + text: str, + *, + normalize: Optional[bool] = None, + truncate: Optional[bool] = None, + model: Optional[str] = None, + ) -> "np.ndarray": """ Generate embeddings for a given text. @@ -926,6 +933,12 @@ async def feature_extraction(self, text: str, *, model: Optional[str] = None) -> The model to use for the conversational task. Can be a model ID hosted on the Hugging Face Hub or a URL to a deployed Inference Endpoint. If not provided, the default recommended conversational model will be used. Defaults to None. + normalize (`bool`, *optional*): + Whether to normalize the embeddings or not. Defaults to None. + Only available on server powered by Text-Embedding-Inference. + truncate (`bool`, *optional*): + Whether to truncate the embeddings or not. Defaults to None. + Only available on server powered by Text-Embedding-Inference. Returns: `np.ndarray`: The embedding representing the input text as a float32 numpy array. @@ -948,7 +961,12 @@ async def feature_extraction(self, text: str, *, model: Optional[str] = None) -> [ 0.28552425, -0.928395 , -1.2077185 , ..., 0.76810825, -2.1069427 , 0.6236161 ]], dtype=float32) ``` """ - response = await self.post(json={"inputs": text}, model=model, task="feature-extraction") + payload = {"inputs": text} + if normalize is not None: + payload["normalize"] = normalize + if truncate is not None: + payload["truncate"] = truncate + response = await self.post(json=payload, model=model, task="feature-extraction") np = _import_numpy() return np.array(_bytes_to_dict(response), dtype="float32") diff --git a/src/huggingface_hub/inference/_generated/types/feature_extraction.py b/src/huggingface_hub/inference/_generated/types/feature_extraction.py index df563e671a..ef87953ac4 100644 --- a/src/huggingface_hub/inference/_generated/types/feature_extraction.py +++ b/src/huggingface_hub/inference/_generated/types/feature_extraction.py @@ -4,16 +4,19 @@ # - 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 Any, Dict, Optional +from typing import List, Optional, Union from .base import BaseInferenceType @dataclass class FeatureExtractionInput(BaseInferenceType): - """Inputs for Text Embedding inference""" + """Feature Extraction Input. + Auto-generated from TEI specs. + For more details, check out + https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tei-import.ts. + """ - inputs: str - """The text to get the embeddings of""" - parameters: Optional[Dict[str, Any]] = None - """Additional inference parameters""" + inputs: Union[List[str], str] + normalize: Optional[bool] = None + truncate: Optional[bool] = None From 62fbfaaa015a199dea68bd41b25ba5343ef05af6 Mon Sep 17 00:00:00 2001 From: Lucain Pouget Date: Fri, 3 May 2024 16:15:28 +0200 Subject: [PATCH 2/2] quality --- src/huggingface_hub/inference/_client.py | 2 +- src/huggingface_hub/inference/_generated/_async_client.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/huggingface_hub/inference/_client.py b/src/huggingface_hub/inference/_client.py index 7122e1149a..1fad9d5837 100644 --- a/src/huggingface_hub/inference/_client.py +++ b/src/huggingface_hub/inference/_client.py @@ -958,7 +958,7 @@ def feature_extraction( [ 0.28552425, -0.928395 , -1.2077185 , ..., 0.76810825, -2.1069427 , 0.6236161 ]], dtype=float32) ``` """ - payload = {"inputs": text} + payload: Dict = {"inputs": text} if normalize is not None: payload["normalize"] = normalize if truncate is not None: diff --git a/src/huggingface_hub/inference/_generated/_async_client.py b/src/huggingface_hub/inference/_generated/_async_client.py index ca8c951319..b1f67a551d 100644 --- a/src/huggingface_hub/inference/_generated/_async_client.py +++ b/src/huggingface_hub/inference/_generated/_async_client.py @@ -961,7 +961,7 @@ async def feature_extraction( [ 0.28552425, -0.928395 , -1.2077185 , ..., 0.76810825, -2.1069427 , 0.6236161 ]], dtype=float32) ``` """ - payload = {"inputs": text} + payload: Dict = {"inputs": text} if normalize is not None: payload["normalize"] = normalize if truncate is not None: