Repository: huggingface/huggingface_hub
Pull request #1607 (2023-08-20 14:04:17)

Add token classification to inference client
**Add  token-classification to HuggingFace🤗 Hub**

References https://github.com/huggingface/huggingface_hub/issues/1539

This is an ongoing list of model tasks to implement in the Hugging Face Hub inference client. Each task is planned to be its own PR. The task for this is [token-classification](https://huggingface.co/tasks/token-classification).

**Key Features**

- Modifies _client.py to call a token classification model.
- Modifies testing and documentation to reflect changes.

HuggingFaceDocBuilderDev: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/huggingface_hub/pr_1607). All of your documentation changes will be reflected on that endpoint.
martinbrose: Thanks for the comments @Wauplin. 

I see if I can find some time this week, but it most likely will be next week.
martinbrose: Hi @Wauplin, I believe I addressed your review points. Please check.

Patch:
diff --git a/docs/source/en/guides/inference.md b/docs/source/en/guides/inference.md
index 5ee7d7a1..70811244 100644
--- a/docs/source/en/guides/inference.md
+++ b/docs/source/en/guides/inference.md
@@ -139,9 +139,9 @@ has a simple API that supports the most common tasks. Here is a list of the curr
 | | [Sentence Similarity](https://huggingface.co/tasks/sentence-similarity)            | ✅ | [`~InferenceClient.sentence_similarity`] |
 | | [Summarization](https://huggingface.co/tasks/summarization)                  | ✅ | [`~InferenceClient.summarization`] |
 | | [Table Question Answering](https://huggingface.co/tasks/table-question-answering)       | |  |
-| | [Text Classification](https://huggingface.co/tasks/text-classification)            | |  |
+| | [Text Classification](https://huggingface.co/tasks/text-classification)            | ✅ | [`~InferenceClient.text_classification`] |
 | | [Text Generation](https://huggingface.co/tasks/text-generation)   | ✅ | [`~InferenceClient.text_generation`] |
-| | [Token Classification](https://huggingface.co/tasks/token-classification)           | |  |
+| | [Token Classification](https://huggingface.co/tasks/token-classification)           | ✅ | [`~InferenceClient.token_classification`] |
 | | [Translation](https://huggingface.co/tasks/translation)       | |  |
 | | [Zero Shot Classification](https://huggingface.co/tasks/zero-shot-image-classification)       | |  |
 | Tabular | [Tabular Classification](https://huggingface.co/tasks/tabular-classification)         | |  |
diff --git a/docs/source/en/package_reference/inference_client.md b/docs/source/en/package_reference/inference_client.md
index 3edeb506..bc14a9e2 100644
--- a/docs/source/en/package_reference/inference_client.md
+++ b/docs/source/en/package_reference/inference_client.md
@@ -49,6 +49,8 @@ For most tasks, the return value has a built-in type (string, list, image...). H
 
 [[autodoc]] huggingface_hub.inference._types.ImageSegmentationOutput
 
+[[autodoc]] huggingface_hub.inference._types.TokenClassificationOutput
+
 ### Text generation types
 
 [`~InferenceClient.text_generation`] task has a greater support than other tasks in `InferenceClient`. In
diff --git a/src/huggingface_hub/inference/_client.py b/src/huggingface_hub/inference/_client.py
index 308e1739..76d13b2c 100644
--- a/src/huggingface_hub/inference/_client.py
+++ b/src/huggingface_hub/inference/_client.py
@@ -82,6 +82,7 @@ from huggingface_hub.inference._types import (
     ConversationalOutput,
     ImageSegmentationOutput,
     ObjectDetectionOutput,
+    TokenClassificationOutput,
 )
 from huggingface_hub.utils import (
     BadRequestError,
@@ -765,6 +766,38 @@ class InferenceClient:
         response = self.post(json=payload, model=model, task="summarization")
         return _bytes_to_dict(response)[0]["summary_text"]
 
+    def text_classification(self, text: str, *, model: Optional[str] = None) -> List[ClassificationOutput]:
+        """
+        Perform sentiment-analysis on the given text.
+
+        Args:
+            text (`str`):
+                A string to be classified.
+            model (`str`, *optional*):
+                The model to use for the text classification 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 text classification model will be used.
+                Defaults to None.
+
+        Returns:
+            `List[Dict]`: a list of dictionaries containing the predicted label and associated probability.
+
+        Raises:
+            [`InferenceTimeoutError`]:
+                If the model is unavailable or the request times out.
+            `HTTPError`:
+                If the request fails with an HTTP error status code other than HTTP 503.
+
+        Example:
+        ```py
+        >>> from huggingface_hub import InferenceClient
+        >>> client = InferenceClient()
+        >>> output = client.text_classification("I like you")
+        [{'label': 'POSITIVE', 'score': 0.9998695850372314}, {'label': 'NEGATIVE', 'score': 0.0001304351753788069}]
+        ```
+        """
+        response = self.post(json={"inputs": text}, model=model, task="text-classification")
+        return _bytes_to_list(response)[0]
+
     @overload
     def text_generation(  # type: ignore
         self,
@@ -1240,6 +1273,53 @@ class InferenceClient:
         """
         return self.post(json={"inputs": text}, model=model, task="text-to-speech")
 
+    def token_classification(self, text: str, *, model: Optional[str] = None) -> List[TokenClassificationOutput]:
+        """
+        Perform token classification on the given text.
+        Usually used for sentence parsing, either grammatical, or Named Entity Recognition (NER) to understand keywords contained within text.
+
+        Args:
+            text (`str`):
+                A string to be classified.
+            model (`str`, *optional*):
+                The model to use for the token classification 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 token classification model will be used.
+                Defaults to None.
+
+        Returns:
+            `List[Dict]`: List of token classification outputs containing the entity group, confidence score, word, start and end index.
+
+        Raises:
+            [`InferenceTimeoutError`]:
+                If the model is unavailable or the request times out.
+            `HTTPError`:
+                If the request fails with an HTTP error status code other than HTTP 503.
+
+        Example:
+        ```py
+        >>> from huggingface_hub import InferenceClient
+        >>> client = InferenceClient()
+        >>> client.token_classification("My name is Sarah Jessica Parker but you can call me Jessica")
+        [{'entity_group': 'PER',
+        'score': 0.9971321225166321,
+        'word': 'Sarah Jessica Parker',
+        'start': 11,
+        'end': 31},
+        {'entity_group': 'PER',
+        'score': 0.9773476123809814,
+        'word': 'Jessica',
+        'start': 52,
+        'end': 59}]
+        ```
+        """
+        payload: Dict[str, Any] = {"inputs": text}
+        response = self.post(
+            json=payload,
+            model=model,
+            task="token-classification",
+        )
+        return _bytes_to_list(response)
+
     def zero_shot_image_classification(
         self, image: ContentT, labels: List[str], *, model: Optional[str] = None
     ) -> List[ClassificationOutput]:
diff --git a/src/huggingface_hub/inference/_generated/_async_client.py b/src/huggingface_hub/inference/_generated/_async_client.py
index 90629a99..e1a9b7e5 100644
--- a/src/huggingface_hub/inference/_generated/_async_client.py
+++ b/src/huggingface_hub/inference/_generated/_async_client.py
@@ -66,6 +66,7 @@ from huggingface_hub.inference._types import (
     ConversationalOutput,
     ImageSegmentationOutput,
     ObjectDetectionOutput,
+    TokenClassificationOutput,
 )
 from huggingface_hub.utils import (
     build_hf_headers,
@@ -772,6 +773,39 @@ class AsyncInferenceClient:
         response = await self.post(json=payload, model=model, task="summarization")
         return _bytes_to_dict(response)[0]["summary_text"]
 
+    async def text_classification(self, text: str, *, model: Optional[str] = None) -> List[ClassificationOutput]:
+        """
+        Perform sentiment-analysis on the given text.
+
+        Args:
+            text (`str`):
+                A string to be classified.
+            model (`str`, *optional*):
+                The model to use for the text classification 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 text classification model will be used.
+                Defaults to None.
+
+        Returns:
+            `List[Dict]`: a list of dictionaries containing the predicted label and associated probability.
+
+        Raises:
+            [`InferenceTimeoutError`]:
+                If the model is unavailable or the request times out.
+            `aiohttp.ClientResponseError`:
+                If the request fails with an HTTP error status code other than HTTP 503.
+
+        Example:
+        ```py
+        # Must be run in an async context
+        >>> from huggingface_hub import AsyncInferenceClient
+        >>> client = AsyncInferenceClient()
+        >>> output = await client.text_classification("I like you")
+        [{'label': 'POSITIVE', 'score': 0.9998695850372314}, {'label': 'NEGATIVE', 'score': 0.0001304351753788069}]
+        ```
+        """
+        response = await self.post(json={"inputs": text}, model=model, task="text-classification")
+        return _bytes_to_list(response)[0]
+
     @overload
     async def text_generation(  # type: ignore
         self,
@@ -1251,6 +1285,54 @@ class AsyncInferenceClient:
         """
         return await self.post(json={"inputs": text}, model=model, task="text-to-speech")
 
+    async def token_classification(self, text: str, *, model: Optional[str] = None) -> List[TokenClassificationOutput]:
+        """
+        Perform token classification on the given text.
+        Usually used for sentence parsing, either grammatical, or Named Entity Recognition (NER) to understand keywords contained within text.
+
+        Args:
+            text (`str`):
+                A string to be classified.
+            model (`str`, *optional*):
+                The model to use for the token classification 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 token classification model will be used.
+                Defaults to None.
+
+        Returns:
+            `List[Dict]`: List of token classification outputs containing the entity group, confidence score, word, start and end index.
+
+        Raises:
+            [`InferenceTimeoutError`]:
+                If the model is unavailable or the request times out.
+            `aiohttp.ClientResponseError`:
+                If the request fails with an HTTP error status code other than HTTP 503.
+
+        Example:
+        ```py
+        # Must be run in an async context
+        >>> from huggingface_hub import AsyncInferenceClient
+        >>> client = AsyncInferenceClient()
+        >>> await client.token_classification("My name is Sarah Jessica Parker but you can call me Jessica")
+        [{'entity_group': 'PER',
+        'score': 0.9971321225166321,
+        'word': 'Sarah Jessica Parker',
+        'start': 11,
+        'end': 31},
+        {'entity_group': 'PER',
+        'score': 0.9773476123809814,
+        'word': 'Jessica',
+        'start': 52,
+        'end': 59}]
+        ```
+        """
+        payload: Dict[str, Any] = {"inputs": text}
+        response = await self.post(
+            json=payload,
+            model=model,
+            task="token-classification",
+        )
+        return _bytes_to_list(response)
+
     async def zero_shot_image_classification(
         self, image: ContentT, labels: List[str], *, model: Optional[str] = None
     ) -> List[ClassificationOutput]:
diff --git a/src/huggingface_hub/inference/_types.py b/src/huggingface_hub/inference/_types.py
index 3af2717a..9ea3d6aa 100644
--- a/src/huggingface_hub/inference/_types.py
+++ b/src/huggingface_hub/inference/_types.py
@@ -98,3 +98,26 @@ class ObjectDetectionOutput(TypedDict):
     label: str
     box: dict
     score: float
+
+
+class TokenClassificationOutput(TypedDict):
+    """Dictionary containing the output of a [`~InferenceClient.token_classification`] task.
+
+    Args:
+        entity_group (`str`):
+            The type for the entity being recognized (model specific).
+        score (`float`):
+            The score of the label predicted by the model.
+        word (`str`):
+            The string that was captured.
+        start (`int`):
+            The offset stringwise where the answer is located. Useful to disambiguate if word occurs multiple times.
+        end (`int`):
+            The offset stringwise where the answer is located. Useful to disambiguate if word occurs multiple times.
+    """
+
+    entity_group: str
+    score: float
+    word: str
+    start: int
+    end: int


Test patch:
diff --git a/tests/cassettes/InferenceClientVCRTest.test_text_classification.yaml b/tests/cassettes/InferenceClientVCRTest.test_text_classification.yaml
new file mode 100644
index 00000000..67e8ac80
--- /dev/null
+++ b/tests/cassettes/InferenceClientVCRTest.test_text_classification.yaml
@@ -0,0 +1,48 @@
+interactions:
+- request:
+    body: '{"inputs": ["I like you"]}'
+    headers:
+      Accept:
+      - '*/*'
+      Accept-Encoding:
+      - gzip, deflate, br
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '41'
+      Content-Type:
+      - application/json
+      X-Amzn-Trace-Id:
+      - b658f44b-c82c-4a0c-9fc1-c287ea0b66d3
+      user-agent:
+      - unknown/None; hf_hub/0.17.0.dev0; python/3.10.12
+    method: POST
+    uri: https://api-inference.huggingface.co/models/distilbert-base-uncased-finetuned-sst-2-english
+  response:
+    body:
+      string: '[[{"label":"POSITIVE","score":0.9998695850372314},{"label":"NEGATIVE","score":0.0001304351753788069}]]'
+    headers:
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '204'
+      Content-Type:
+      - application/json
+      Date:
+      - Sun, 20 Aug 2023 11:48:55 GMT
+      access-control-allow-credentials:
+      - 'true'
+      vary:
+      - Origin, Access-Control-Request-Method, Access-Control-Request-Headers
+      x-compute-time:
+      - '0.033'
+      x-compute-type:
+      - cache
+      x-request-id:
+      - MiuTWky1u3OlV7JlitniT
+      x-sha:
+      - 3d65bad49c7ba6f71920504507a8927f4b9db6c0
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/cassettes/InferenceClientVCRTest.test_token_classification.yaml b/tests/cassettes/InferenceClientVCRTest.test_token_classification.yaml
new file mode 100644
index 00000000..4db6c7e5
--- /dev/null
+++ b/tests/cassettes/InferenceClientVCRTest.test_token_classification.yaml
@@ -0,0 +1,49 @@
+interactions:
+- request:
+    body: '{"inputs": "My name is Sarah Jessica Parker but you can call me Jessica"}'
+    headers:
+      Accept:
+      - '*/*'
+      Accept-Encoding:
+      - gzip, deflate, br
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '73'
+      Content-Type:
+      - application/json
+      X-Amzn-Trace-Id:
+      - d19b39c0-95d3-4879-b316-3b4dbea8d77e
+      user-agent:
+      - unknown/None; hf_hub/0.17.0.dev0; python/3.10.12
+    method: POST
+    uri: https://api-inference.huggingface.co/models/dbmdz/bert-large-cased-finetuned-conll03-english
+  response:
+    body:
+      string: '[{"entity_group":"PER","score":0.9991335868835449,"word":"Sarah Jessica
+        Parker","start":11,"end":31},{"entity_group":"PER","score":0.9979913234710693,"word":"Jessica","start":52,"end":59}]'
+    headers:
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '188'
+      Content-Type:
+      - application/json
+      Date:
+      - Sun, 20 Aug 2023 14:01:19 GMT
+      access-control-allow-credentials:
+      - 'true'
+      vary:
+      - Origin, Access-Control-Request-Method, Access-Control-Request-Headers
+      x-compute-time:
+      - '0.070'
+      x-compute-type:
+      - cache
+      x-request-id:
+      - 2trJlwwa5RQvr-9LZEsq3
+      x-sha:
+      - f2482bf01f5da0f0eb8e183ffd8cc3885aa90b14
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/test_inference_client.py b/tests/test_inference_client.py
index bb0f6366..5a526cf3 100644
--- a/tests/test_inference_client.py
+++ b/tests/test_inference_client.py
@@ -200,6 +200,14 @@ class InferenceClientVCRTest(InferenceClientTest):
             " surpassed the Washington Monument to become the tallest man-made structure in the world.",
         )
 
+    def test_text_classification(self) -> None:
+        output = self.client.text_classification("I like you")
+        self.assertIsInstance(output, list)
+        self.assertEqual(len(output), 2)
+        for item in output:
+            self.assertIsInstance(item["score"], float)
+            self.assertIsInstance(item["label"], str)
+
     def test_text_generation(self) -> None:
         """Tested separately in `test_inference_text_generation.py`."""
 
@@ -219,6 +227,20 @@ class InferenceClientVCRTest(InferenceClientTest):
         audio = self.client.text_to_speech("Hello world")
         self.assertIsInstance(audio, bytes)
 
+    def test_token_classification(self) -> None:
+        model = "dbmdz/bert-large-cased-finetuned-conll03-english"
+        output = self.client.token_classification(
+            "My name is Sarah Jessica Parker but you can call me Jessica", model=model
+        )
+        self.assertIsInstance(output, list)
+        self.assertGreater(len(output), 0)
+        for item in output:
+            self.assertIsInstance(item["entity_group"], str)
+            self.assertIsInstance(item["score"], float)
+            self.assertIsInstance(item["word"], str)
+            self.assertIsInstance(item["start"], int)
+            self.assertIsInstance(item["end"], int)
+
     def test_zero_shot_image_classification(self) -> None:
         output = self.client.zero_shot_image_classification(self.image_file, ["tree", "woman", "cat"])
         self.assertIsInstance(output, list)