fix tests

huggingface · Nov 20, 2024 · 690decf · 690decf
1 parent b66ba40
commit 690decf
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 10 deletions.
diff --git a/src/huggingface_hub/inference/_client.py b/src/huggingface_hub/inference/_client.py
@@ -979,7 +979,7 @@ def document_question_answering(
         >>> from huggingface_hub import InferenceClient
         >>> client = InferenceClient()
         >>> client.document_question_answering(image="https://huggingface.co/spaces/impira/docquery/resolve/2359223c1837a7587402bda0f2643382a6eefeab/invoice.png", question="What is the invoice number?")
-        [DocumentQuestionAnsweringOutputElement(answer='us-001', end=16, score=0.9999666213989258, start=16, words=None)]
+        [DocumentQuestionAnsweringOutputElement(answer='us-001', end=16, score=0.9999666213989258, start=16)]
         ```
         """
         inputs: Dict[str, Any] = {"question": question, "image": _b64_encode(image)}

diff --git a/src/huggingface_hub/inference/_generated/_async_client.py b/src/huggingface_hub/inference/_generated/_async_client.py
@@ -1022,7 +1022,7 @@ async def document_question_answering(
         >>> from huggingface_hub import AsyncInferenceClient
         >>> client = AsyncInferenceClient()
         >>> await client.document_question_answering(image="https://huggingface.co/spaces/impira/docquery/resolve/2359223c1837a7587402bda0f2643382a6eefeab/invoice.png", question="What is the invoice number?")
-        [DocumentQuestionAnsweringOutputElement(answer='us-001', end=16, score=0.9999666213989258, start=16, words=None)]
+        [DocumentQuestionAnsweringOutputElement(answer='us-001', end=16, score=0.9999666213989258, start=16)]
         ```
         """
         inputs: Dict[str, Any] = {"question": question, "image": _b64_encode(image)}

diff --git a/tests/test_inference_client.py b/tests/test_inference_client.py
@@ -390,7 +390,10 @@ def test_document_question_answering(self) -> None:
             output,
             [
                 DocumentQuestionAnsweringOutputElement(
-                    answer="$1,000,000,000", end=None, score=None, start=None, words=None
+                    answer="$1,000,000,000",
+                    end=None,
+                    score=None,
+                    start=None,
                 )
             ],
         )
@@ -619,21 +622,21 @@ def test_token_classification(self) -> None:
         output = self.client.token_classification("My name is Sarah Jessica Parker but you can call me Jessica")
         assert output == [
             TokenClassificationOutputElement(
-                label=None, score=0.9991335868835449, end=31, entity_group="PER", start=11, word="Sarah Jessica Parker"
+                score=0.9991335868835449, end=31, entity_group="PER", start=11, word="Sarah Jessica Parker"
             ),
             TokenClassificationOutputElement(
-                label=None, score=0.9979913234710693, end=59, entity_group="PER", start=52, word="Jessica"
+                score=0.9979913234710693, end=59, entity_group="PER", start=52, word="Jessica"
             ),
         ]
 
     def test_visual_question_answering(self) -> None:
         output = self.client.visual_question_answering(self.image_file, "Who's in the picture?")
         assert output == [
-            VisualQuestionAnsweringOutputElement(label=None, score=0.9386941194534302, answer="woman"),
-            VisualQuestionAnsweringOutputElement(label=None, score=0.34311845898628235, answer="girl"),
-            VisualQuestionAnsweringOutputElement(label=None, score=0.08407749235630035, answer="lady"),
-            VisualQuestionAnsweringOutputElement(label=None, score=0.0507517009973526, answer="female"),
-            VisualQuestionAnsweringOutputElement(label=None, score=0.01777094043791294, answer="man"),
+            VisualQuestionAnsweringOutputElement(score=0.9386941194534302, answer="woman"),
+            VisualQuestionAnsweringOutputElement(score=0.34311845898628235, answer="girl"),
+            VisualQuestionAnsweringOutputElement(score=0.08407749235630035, answer="lady"),
+            VisualQuestionAnsweringOutputElement(score=0.0507517009973526, answer="female"),
+            VisualQuestionAnsweringOutputElement(score=0.01777094043791294, answer="man"),
         ]
 
     @expect_deprecation("zero_shot_classification")