Generic pipeline is all you need (#371)

osanseviero · web-flow · commit dea8c4acc91b · 2021-09-27T08:28:53.000+02:00
diff --git a/.github/workflows/python-api-generic.yaml b/.github/workflows/python-api-generic.yaml
@@ -0,0 +1,29 @@
+# TODO: Merge this with allenNLP to have a single workflow for all docker images.
+name: generic-docker
+
+on:
+  pull_request:
+    paths:
+      - "api-inference-community/docker_images/generic/**"
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: "3.8"
+      - name: Checkout
+        uses: actions/checkout@v2
+      - name: Set up QEMU
+        uses: docker/setup-qemu-action@v1
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v1
+      - name: Install dependencies
+        working-directory: api-inference-community
+        run: |
+          pip install --upgrade pip
+          pip install pytest pillow httpx
+          pip install -e .
+      - run: RUN_DOCKER_TESTS=1 pytest -sv tests/test_dockers.py::DockerImageTests::test_generic
+        working-directory: api-inference-community
diff --git a/api-inference-community/docker_images/generic/app/main.py b/api-inference-community/docker_images/generic/app/main.py
@@ -1,7 +1,6 @@
 import functools
 import logging
 import os
-from typing import Dict, Type
 
 from api_inference_community.routes import pipeline_route, status_ok
 from app.pipelines import Pipeline
@@ -18,24 +17,10 @@
 logger = logging.getLogger(__name__)
 
 
-ALLOWED_TASKS: Dict[str, Type[Pipeline]] = {
-    "audio-to-audio": Pipeline,
-    "automatic-speech-recognition": Pipeline,
-    "feature-extraction": Pipeline,
-    "image-classification": Pipeline,
-    "structured-data-classification": Pipeline,
-    "text-to-image": Pipeline,
-    "token-classification": Pipeline,
-}
-
-
 @functools.lru_cache()
 def get_pipeline() -> Pipeline:
-    task = os.environ["TASK"]
     model_id = os.environ["MODEL_ID"]
-    if task not in ALLOWED_TASKS:
-        raise EnvironmentError(f"{task} is not a valid pipeline for model : {model_id}")
-    return ALLOWED_TASKS[task](model_id)
+    return Pipeline(model_id)
 
 
 routes = [
diff --git a/api-inference-community/docker_images/generic/tests/test_api.py b/api-inference-community/docker_images/generic/tests/test_api.py
@@ -1,8 +1,4 @@
-import os
 from typing import Dict, List
-from unittest import TestCase, skipIf
-
-from app.main import ALLOWED_TASKS, get_pipeline
 
 
 # Must contain at least one example of each implemented pipeline
@@ -14,39 +10,8 @@
     # This is very slow the first time as fasttext model is large.
     "feature-extraction": ["osanseviero/fasttext_english"],
     "image-classification": ["osanseviero/fastai_cat_vs_dog"],
+    "structured-data-classification": ["osanseviero/wine-quality"],
+    "text-classification": ["osanseviero/fasttext_nearest"],
     "text-to-image": ["osanseviero/BigGAN-deep-128"],
     "token-classification": ["osanseviero/en_core_web_sm"],
-    "structured-data-classification": ["osanseviero/wine-quality"],
-}
-
-ALL_TASKS = {
-    "audio-to-audio",
-    "automatic-speech-recognition",
-    "feature-extraction",
-    "image-classification",
-    "question-answering",
-    "sentence-similarity",
-    "structure-data-classification",
-    "text-to-speech",
-    "token-classification",
 }
-
-
-class PipelineTestCase(TestCase):
-    @skipIf(
-        os.path.dirname(os.path.dirname(__file__)).endswith("common"),
-        "common is a special case",
-    )
-    def test_has_at_least_one_task_enabled(self):
-        self.assertGreater(
-            len(ALLOWED_TASKS.keys()), 0, "You need to implement at least one task"
-        )
-
-    def test_unsupported_tasks(self):
-        unsupported_tasks = ALL_TASKS - ALLOWED_TASKS.keys()
-        for unsupported_task in unsupported_tasks:
-            with self.subTest(msg=unsupported_task, task=unsupported_task):
-                os.environ["TASK"] = unsupported_task
-                os.environ["MODEL_ID"] = "XX"
-                with self.assertRaises(EnvironmentError):
-                    get_pipeline()
diff --git a/api-inference-community/docker_images/generic/tests/test_api_audio_to_audio.py b/api-inference-community/docker_images/generic/tests/test_api_audio_to_audio.py
@@ -4,14 +4,13 @@
 from unittest import TestCase, skipIf
 
 from api_inference_community.validation import ffmpeg_read
-from app.main import ALLOWED_TASKS
 from parameterized import parameterized_class
 from starlette.testclient import TestClient
 from tests.test_api import TESTABLE_MODELS
 
 
 @skipIf(
-    "audio-to-audio" not in ALLOWED_TASKS,
+    "audio-to-audio" not in TESTABLE_MODELS,
     "audio-to-audio not implemented",
 )
 @parameterized_class(
diff --git a/api-inference-community/docker_images/generic/tests/test_api_automatic_speech_recognition.py b/api-inference-community/docker_images/generic/tests/test_api_automatic_speech_recognition.py
@@ -2,14 +2,13 @@
 import os
 from unittest import TestCase, skipIf
 
-from app.main import ALLOWED_TASKS
 from parameterized import parameterized_class
 from starlette.testclient import TestClient
 from tests.test_api import TESTABLE_MODELS
 
 
 @skipIf(
-    "automatic-speech-recognition" not in ALLOWED_TASKS,
+    "automatic-speech-recognition" not in TESTABLE_MODELS,
     "automatic-speech-recognition not implemented",
 )
 @parameterized_class(
diff --git a/api-inference-community/docker_images/generic/tests/test_api_feature_extraction.py b/api-inference-community/docker_images/generic/tests/test_api_feature_extraction.py
@@ -2,14 +2,13 @@
 import os
 from unittest import TestCase, skipIf
 
-from app.main import ALLOWED_TASKS
 from parameterized import parameterized_class
 from starlette.testclient import TestClient
 from tests.test_api import TESTABLE_MODELS
 
 
 @skipIf(
-    "feature-extraction" not in ALLOWED_TASKS,
+    "feature-extraction" not in TESTABLE_MODELS,
     "feature-extraction not implemented",
 )
 @parameterized_class(
diff --git a/api-inference-community/docker_images/generic/tests/test_api_image_classification.py b/api-inference-community/docker_images/generic/tests/test_api_image_classification.py
@@ -2,14 +2,13 @@
 import os
 from unittest import TestCase, skipIf
 
-from app.main import ALLOWED_TASKS
 from parameterized import parameterized_class
 from starlette.testclient import TestClient
 from tests.test_api import TESTABLE_MODELS
 
 
 @skipIf(
-    "image-classification" not in ALLOWED_TASKS,
+    "image-classification" not in TESTABLE_MODELS,
     "image-classification not implemented",
 )
 @parameterized_class(
diff --git a/api-inference-community/docker_images/generic/tests/test_api_question_answering.py b/api-inference-community/docker_images/generic/tests/test_api_question_answering.py
@@ -2,14 +2,13 @@
 import os
 from unittest import TestCase, skipIf
 
-from app.main import ALLOWED_TASKS
 from starlette.testclient import TestClient
 from tests.test_api import TESTABLE_MODELS
 
 
 @skipIf(
-    "question-answering" not in ALLOWED_TASKS,
-    "question-answering not implemented",
+    "text-to-speech" not in TESTABLE_MODELS,
+    "text-to-speech not implemented",
 )
 class QuestionAnsweringTestCase(TestCase):
     def setUp(self):
diff --git a/api-inference-community/docker_images/generic/tests/test_api_sentence_similarity.py b/api-inference-community/docker_images/generic/tests/test_api_sentence_similarity.py
@@ -2,13 +2,12 @@
 import os
 from unittest import TestCase, skipIf
 
-from app.main import ALLOWED_TASKS
 from starlette.testclient import TestClient
 from tests.test_api import TESTABLE_MODELS
 
 
 @skipIf(
-    "sentence-similarity" not in ALLOWED_TASKS,
+    "sentence-similarity" not in TESTABLE_MODELS,
     "sentence-similarity not implemented",
 )
 class SentenceSimilarityTestCase(TestCase):
diff --git a/api-inference-community/docker_images/generic/tests/test_api_speech_segmentation.py b/api-inference-community/docker_images/generic/tests/test_api_speech_segmentation.py
@@ -2,14 +2,13 @@
 import os
 from unittest import TestCase, skipIf
 
-from app.main import ALLOWED_TASKS
 from starlette.testclient import TestClient
 from tests.test_api import TESTABLE_MODELS
 
 
 @skipIf(
-    "speech-segmentation" not in ALLOWED_TASKS,
-    "speech-segmentation not implemented",
+    "text-to-speech" not in TESTABLE_MODELS,
+    "text-to-speech not implemented",
 )
 class SpeechSegmentationTestCase(TestCase):
     def setUp(self):
diff --git a/api-inference-community/docker_images/generic/tests/test_api_structured_data_classification.py b/api-inference-community/docker_images/generic/tests/test_api_structured_data_classification.py
@@ -2,14 +2,13 @@
 import os
 from unittest import TestCase, skipIf
 
-from app.main import ALLOWED_TASKS
 from parameterized import parameterized_class
 from starlette.testclient import TestClient
 from tests.test_api import TESTABLE_MODELS
 
 
 @skipIf(
-    "structured-data-classification" not in ALLOWED_TASKS,
+    "structured-data-classification" not in TESTABLE_MODELS,
     "structured-data-classification not implemented",
 )
 @parameterized_class(
diff --git a/api-inference-community/docker_images/generic/tests/test_api_text_classification.py b/api-inference-community/docker_images/generic/tests/test_api_text_classification.py
@@ -0,0 +1,93 @@
+import json
+import os
+from unittest import TestCase, skipIf
+
+from parameterized import parameterized_class
+from starlette.testclient import TestClient
+from tests.test_api import TESTABLE_MODELS
+
+
+@skipIf(
+    "text-classification" not in TESTABLE_MODELS,
+    "text-classification not implemented",
+)
+@parameterized_class(
+    [{"model_id": model_id} for model_id in TESTABLE_MODELS["text-classification"]]
+)
+class TextClassificationTestCase(TestCase):
+    def setUp(self):
+        self.old_model_id = os.getenv("MODEL_ID")
+        self.old_task = os.getenv("TASK")
+        os.environ["MODEL_ID"] = self.model_id
+        os.environ["TASK"] = "token-classification"
+
+        from app.main import app, get_pipeline
+
+        get_pipeline.cache_clear()
+
+        self.app = app
+
+    @classmethod
+    def setUpClass(cls):
+        from app.main import get_pipeline
+
+        get_pipeline.cache_clear()
+
+    def tearDown(self):
+        if self.old_model_id is not None:
+            os.environ["MODEL_ID"] = self.old_model_id
+        else:
+            del os.environ["MODEL_ID"]
+        if self.old_task is not None:
+            os.environ["TASK"] = self.old_task
+        else:
+            del os.environ["TASK"]
+
+    def test_simple(self):
+        inputs = "It is a beautiful day outside"
+
+        with TestClient(self.app) as client:
+            response = client.post("/", json={"inputs": inputs})
+        self.assertEqual(
+            response.status_code,
+            200,
+        )
+        content = json.loads(response.content)
+        self.assertEqual(type(content), list)
+        self.assertEqual(len(content), 1)
+        self.assertEqual(type(content[0]), list)
+        self.assertEqual(
+            set(k for el in content[0] for k in el.keys()),
+            {"label", "score"},
+        )
+
+        with TestClient(self.app) as client:
+            response = client.post("/", json=inputs)
+
+        self.assertEqual(
+            response.status_code,
+            200,
+        )
+        content = json.loads(response.content)
+        self.assertEqual(type(content), list)
+        self.assertEqual(len(content), 1)
+        self.assertEqual(type(content[0]), list)
+        self.assertEqual(
+            set(k for el in content[0] for k in el.keys()),
+            {"label", "score"},
+        )
+
+    """
+    def test_malformed_question(self):
+        with TestClient(self.app) as client:
+            response = client.post("/", data=b"\xc3\x28")
+
+        self.assertEqual(
+            response.status_code,
+            400,
+        )
+        self.assertEqual(
+            response.content,
+            b'{"error":"\'utf-8\' codec can\'t decode byte 0xc3 in position 0: invalid continuation byte"}',
+        )
+    """
diff --git a/api-inference-community/docker_images/generic/tests/test_api_text_to_image.py b/api-inference-community/docker_images/generic/tests/test_api_text_to_image.py
@@ -3,14 +3,13 @@
 from unittest import TestCase, skipIf
 
 import PIL
-from app.main import ALLOWED_TASKS
 from parameterized import parameterized_class
 from starlette.testclient import TestClient
 from tests.test_api import TESTABLE_MODELS
 
 
 @skipIf(
-    "text-to-image" not in ALLOWED_TASKS,
+    "text-to-image" not in TESTABLE_MODELS,
     "text-to-image not implemented",
 )
 @parameterized_class(
diff --git a/api-inference-community/docker_images/generic/tests/test_api_text_to_speech.py b/api-inference-community/docker_images/generic/tests/test_api_text_to_speech.py
@@ -2,13 +2,12 @@
 from unittest import TestCase, skipIf
 
 from api_inference_community.validation import ffmpeg_read
-from app.main import ALLOWED_TASKS
 from starlette.testclient import TestClient
 from tests.test_api import TESTABLE_MODELS
 
 
 @skipIf(
-    "text-to-speech" not in ALLOWED_TASKS,
+    "text-to-speech" not in TESTABLE_MODELS,
     "text-to-speech not implemented",
 )
 class TextToSpeechTestCase(TestCase):
diff --git a/api-inference-community/docker_images/generic/tests/test_api_token_classification.py b/api-inference-community/docker_images/generic/tests/test_api_token_classification.py
@@ -2,14 +2,13 @@
 import os
 from unittest import TestCase, skipIf
 
-from app.main import ALLOWED_TASKS
 from parameterized import parameterized_class
 from starlette.testclient import TestClient
 from tests.test_api import TESTABLE_MODELS
 
 
 @skipIf(
-    "token-classification" not in ALLOWED_TASKS,
+    "token-classification" not in TESTABLE_MODELS,
     "token-classification not implemented",
 )
 @parameterized_class(
