Repository: huggingface/huggingface_hub
Pull request #1613 (2023-08-21 15:40:41)

Add fill mask to inference client
**Add fill-mask 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 [fill-mask](https://huggingface.co/tasks/fill-mask).

**Key Features**

- Modifies _client.py to call a fill-mask 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_1613). All of your documentation changes will be reflected on that endpoint.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1613?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) Report
Patch coverage: **`76.92%`** and project coverage change: **`-0.46%`** :warning:
> Comparison is base [(`e328730`)](https://app.codecov.io/gh/huggingface/huggingface_hub/commit/e328730821be171d42e39155d5ab08b52e160a9d?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) 82.26% compared to head [(`f735f4f`)](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1613?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) 81.81%.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #1613      +/-   ##
==========================================
- Coverage   82.26%   81.81%   -0.46%     
==========================================
  Files          60       60              
  Lines        6777     6790      +13     
==========================================
- Hits         5575     5555      -20     
- Misses       1202     1235      +33     
```


| [Files Changed](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1613?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface) | Coverage Δ | |
|---|---|---|
| [...gingface\_hub/inference/\_generated/\_async\_client.py](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1613?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface#diff-c3JjL2h1Z2dpbmdmYWNlX2h1Yi9pbmZlcmVuY2UvX2dlbmVyYXRlZC9fYXN5bmNfY2xpZW50LnB5) | `58.37% <25.00%> (-0.62%)` | :arrow_down: |
| [src/huggingface\_hub/inference/\_client.py](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1613?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface#diff-c3JjL2h1Z2dpbmdmYWNlX2h1Yi9pbmZlcmVuY2UvX2NsaWVudC5weQ==) | `79.51% <100.00%> (+0.40%)` | :arrow_up: |
| [src/huggingface\_hub/inference/\_types.py](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1613?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface#diff-c3JjL2h1Z2dpbmdmYWNlX2h1Yi9pbmZlcmVuY2UvX3R5cGVzLnB5) | `96.15% <100.00%> (+0.91%)` | :arrow_up: |

... and [5 files with indirect coverage changes](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1613/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/huggingface/huggingface_hub/pull/1613?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=huggingface).


Patch:
diff --git a/docs/source/en/guides/inference.md b/docs/source/en/guides/inference.md
index e0e42076..c16d6bbd 100644
--- a/docs/source/en/guides/inference.md
+++ b/docs/source/en/guides/inference.md
@@ -134,7 +134,7 @@ has a simple API that supports the most common tasks. Here is a list of the curr
 | | [Visual Question Answering](https://huggingface.co/tasks/visual-question-answering)      |              |  |
 | NLP | [Conversational](https://huggingface.co/tasks/conversational)                 | ✅ | [`~InferenceClient.conversational`] |
 | | [Feature Extraction](https://huggingface.co/tasks/feature-extraction)             | ✅ | [`~InferenceClient.feature_extraction`] |
-| | [Fill Mask](https://huggingface.co/tasks/fill-mask)                      |              |  |
+| | [Fill Mask](https://huggingface.co/tasks/fill-mask)                      | ✅ | [`~InferenceClient.fill_mask`] |
 | | [Question Answering](https://huggingface.co/tasks/question-answering)             | ✅ | [`~InferenceClient.question_answering`]
 | | [Sentence Similarity](https://huggingface.co/tasks/sentence-similarity)            | ✅ | [`~InferenceClient.sentence_similarity`] |
 | | [Summarization](https://huggingface.co/tasks/summarization)                  | ✅ | [`~InferenceClient.summarization`] |
diff --git a/src/huggingface_hub/inference/_client.py b/src/huggingface_hub/inference/_client.py
index 0bc301a7..c1073e47 100644
--- a/src/huggingface_hub/inference/_client.py
+++ b/src/huggingface_hub/inference/_client.py
@@ -80,6 +80,7 @@ from huggingface_hub.inference._text_generation import (
 from huggingface_hub.inference._types import (
     ClassificationOutput,
     ConversationalOutput,
+    FillMaskOutput,
     ImageSegmentationOutput,
     ObjectDetectionOutput,
     QuestionAnsweringOutput,
@@ -423,6 +424,46 @@ class InferenceClient:
         np = _import_numpy()
         return np.array(_bytes_to_dict(response)[0], dtype="float32")
 
+    def fill_mask(self, text: str, *, model: Optional[str] = None) -> List[FillMaskOutput]:
+        """
+        Fill in a hole with a missing word (token to be precise).
+
+        Args:
+            text (`str`):
+                a string to be filled from, must contain the [MASK] token (check model card for exact name of the mask).
+            model (`str`, *optional*):
+                The model to use for the fill mask 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 fill mask model will be used.
+                Defaults to None.
+
+        Returns:
+            `List[Dict]`: a list of fill mask output dictionaries containing the predicted label, associated
+            probability, token reference, and completed text.
+
+        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.fill_mask("The goal of life is <mask>.")
+        [{'score': 0.06897063553333282,
+        'token': 11098,
+        'token_str': ' happiness',
+        'sequence': 'The goal of life is happiness.'},
+        {'score': 0.06554922461509705,
+        'token': 45075,
+        'token_str': ' immortality',
+        'sequence': 'The goal of life is immortality.'}]
+        ```
+        """
+        response = self.post(json={"inputs": text}, model=model, task="fill-mask")
+        return _bytes_to_list(response)
+
     def image_classification(
         self,
         image: ContentT,
diff --git a/src/huggingface_hub/inference/_generated/_async_client.py b/src/huggingface_hub/inference/_generated/_async_client.py
index 14d4fe16..ec1204a7 100644
--- a/src/huggingface_hub/inference/_generated/_async_client.py
+++ b/src/huggingface_hub/inference/_generated/_async_client.py
@@ -64,6 +64,7 @@ from huggingface_hub.inference._text_generation import (
 from huggingface_hub.inference._types import (
     ClassificationOutput,
     ConversationalOutput,
+    FillMaskOutput,
     ImageSegmentationOutput,
     ObjectDetectionOutput,
     QuestionAnsweringOutput,
@@ -423,6 +424,47 @@ class AsyncInferenceClient:
         np = _import_numpy()
         return np.array(_bytes_to_dict(response)[0], dtype="float32")
 
+    async def fill_mask(self, text: str, *, model: Optional[str] = None) -> List[FillMaskOutput]:
+        """
+        Fill in a hole with a missing word (token to be precise).
+
+        Args:
+            text (`str`):
+                a string to be filled from, must contain the [MASK] token (check model card for exact name of the mask).
+            model (`str`, *optional*):
+                The model to use for the fill mask 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 fill mask model will be used.
+                Defaults to None.
+
+        Returns:
+            `List[Dict]`: a list of fill mask output dictionaries containing the predicted label, associated
+            probability, token reference, and completed text.
+
+        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.fill_mask("The goal of life is <mask>.")
+        [{'score': 0.06897063553333282,
+        'token': 11098,
+        'token_str': ' happiness',
+        'sequence': 'The goal of life is happiness.'},
+        {'score': 0.06554922461509705,
+        'token': 45075,
+        'token_str': ' immortality',
+        'sequence': 'The goal of life is immortality.'}]
+        ```
+        """
+        response = await self.post(json={"inputs": text}, model=model, task="fill-mask")
+        return _bytes_to_list(response)
+
     async def image_classification(
         self,
         image: ContentT,
diff --git a/src/huggingface_hub/inference/_types.py b/src/huggingface_hub/inference/_types.py
index e899ff31..fb5537b9 100644
--- a/src/huggingface_hub/inference/_types.py
+++ b/src/huggingface_hub/inference/_types.py
@@ -64,6 +64,26 @@ class ConversationalOutput(TypedDict):
     warnings: List[str]
 
 
+class FillMaskOutput(TypedDict):
+    """Dictionary containing information about a [`~InferenceClient.fill_mask`] task.
+
+    Args:
+        score (`float`):
+            The probability of the token.
+        token (`int`):
+            The id of the token.
+        token_str (`str`):
+            The string representation of the token.
+        sequence (`str`):
+            The actual sequence of tokens that ran against the model (may contain special tokens).
+    """
+
+    score: float
+    token: int
+    token_str: str
+    sequence: str
+
+
 class ImageSegmentationOutput(TypedDict):
     """Dictionary containing information about a [`~InferenceClient.image_segmentation`] task. In practice, image segmentation returns a
     list of `ImageSegmentationOutput` with 1 item per mask.


Test patch:
diff --git a/tests/cassettes/InferenceClientVCRTest.test_fill_mask.yaml b/tests/cassettes/InferenceClientVCRTest.test_fill_mask.yaml
new file mode 100644
index 00000000..b12dd3e9
--- /dev/null
+++ b/tests/cassettes/InferenceClientVCRTest.test_fill_mask.yaml
@@ -0,0 +1,53 @@
+interactions:
+- request:
+    body: '{"inputs": "The goal of life is <mask>."}'
+    headers:
+      Accept:
+      - '*/*'
+      Accept-Encoding:
+      - gzip, deflate, br
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '41'
+      Content-Type:
+      - application/json
+      X-Amzn-Trace-Id:
+      - fcaea546-88d0-47b5-a5a1-ef49198074c8
+      user-agent:
+      - unknown/None; hf_hub/0.17.0.dev0; python/3.10.12; torch/2.0.0.post101
+    method: POST
+    uri: https://api-inference.huggingface.co/models/distilroberta-base
+  response:
+    body:
+      string: '[{"score":0.06897063553333282,"token":11098,"token_str":" happiness","sequence":"The
+        goal of life is happiness."},{"score":0.06554922461509705,"token":45075,"token_str":"
+        immortality","sequence":"The goal of life is immortality."},{"score":0.0323575921356678,"token":14314,"token_str":"
+        yours","sequence":"The goal of life is yours."},{"score":0.02431388944387436,"token":22211,"token_str":"
+        liberation","sequence":"The goal of life is liberation."},{"score":0.023767812177538872,"token":25342,"token_str":"
+        simplicity","sequence":"The goal of life is simplicity."}]'
+    headers:
+      Connection:
+      - keep-alive
+      Content-Length:
+      - '566'
+      Content-Type:
+      - application/json
+      Date:
+      - Mon, 21 Aug 2023 15:37:19 GMT
+      access-control-allow-credentials:
+      - 'true'
+      vary:
+      - Origin, Access-Control-Request-Method, Access-Control-Request-Headers
+      x-compute-time:
+      - '0.036'
+      x-compute-type:
+      - cache
+      x-request-id:
+      - luPMOhNSkCeOLo7sbj3GE
+      x-sha:
+      - d5411c3ee9e1793fd9ef58390b40a80a4c10df32
+    status:
+      code: 200
+      message: OK
+version: 1
diff --git a/tests/test_inference_client.py b/tests/test_inference_client.py
index 3bdcddd3..56d9860c 100644
--- a/tests/test_inference_client.py
+++ b/tests/test_inference_client.py
@@ -128,6 +128,17 @@ class InferenceClientVCRTest(InferenceClientTest):
         self.assertIsInstance(embedding, np.ndarray)
         self.assertEqual(embedding.shape, (8, 768))
 
+    def test_fill_mask(self) -> None:
+        model = "distilroberta-base"
+        output = self.client.fill_mask("The goal of life is <mask>.", model=model)
+        self.assertIsInstance(output, list)
+        self.assertEqual(len(output[0]), 4)
+        self.assertIsInstance(output[0], dict)
+        self.assertEqual(
+            set(k for el in output for k in el.keys()),
+            {"score", "sequence", "token", "token_str"},
+        )
+
     def test_image_classification(self) -> None:
         output = self.client.image_classification(self.image_file)
         self.assertIsInstance(output, list)