diff --git a/.github/workflows/python-api-thirdai-cd.yaml b/.github/workflows/python-api-thirdai-cd.yaml new file mode 100644 index 00000000..8d59f931 --- /dev/null +++ b/.github/workflows/python-api-thirdai-cd.yaml @@ -0,0 +1,47 @@ +name: thirdai-docker-cd +on: + push: + branches: + - main + paths: + - "docker_images/thirdai/**" +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 + run: | + pip install --upgrade pip + pip install awscli + - uses: tailscale/github-action@v1 + with: + authkey: ${{ secrets.TAILSCALE_AUTHKEY }} + - name: Update upstream + env: + AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} + AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} + AWS_DEFAULT_REGION: ${{ secrets.AWS_DEFAULT_REGION }} + DEFAULT_HOSTNAME: ${{ secrets.DEFAULT_HOSTNAME }} + REGISTRY_USERNAME: ${{ secrets.REGISTRY_USERNAME }} + REGISTRY_PASSWORD: ${{ secrets.REGISTRY_PASSWORD }} + run: | + python build_docker.py thirdai --out out.txt + - name: Deploy on API + run: | + # Load the tags into the env + cat out.txt >> $GITHUB_ENV + export $(xargs < out.txt) + echo ${THIRDAI_CPU_TAG} + # Weird single quote escape mechanism because string interpolation does + # not work on single quote in bash + curl -H "Authorization: Bearer ${{ secrets.API_GITHUB_TOKEN }}" https://api.github.com/repos/huggingface/api-inference/actions/workflows/update_community.yaml/dispatches -d '{"ref":"main","inputs":{"framework":"THIRDAI","tag": "'"${THIRDAI_CPU_TAG}"'"}}' diff --git a/.github/workflows/python-api-thirdai.yaml b/.github/workflows/python-api-thirdai.yaml new file mode 100644 index 00000000..ca507771 --- /dev/null +++ b/.github/workflows/python-api-thirdai.yaml @@ -0,0 +1,26 @@ +name: thirdai-docker + +on: + pull_request: + paths: + - "docker_images/thirdai/**" +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 + 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_thirdai diff --git a/docker_images/thirdai/Dockerfile b/docker_images/thirdai/Dockerfile new file mode 100644 index 00000000..bbe7b680 --- /dev/null +++ b/docker_images/thirdai/Dockerfile @@ -0,0 +1,29 @@ +FROM --platform=linux/x86_64 tiangolo/uvicorn-gunicorn:python3.10 +LABEL maintainer="David Torres " + +# Add any system dependency here +# RUN apt-get update -y && apt-get install libXXX -y + +COPY ./requirements.txt /app +RUN pip install --no-cache-dir -r requirements.txt +COPY ./prestart.sh /app/ + + +# Most DL models are quite large in terms of memory, using workers is a HUGE +# slowdown because of the fork and GIL with python. +# Using multiple pods seems like a better default strategy. +# Feel free to override if it does not make sense for your library. +ARG max_workers=1 +ENV MAX_WORKERS=$max_workers +ENV HUGGINGFACE_HUB_CACHE=/data + +# Necessary on GPU environment docker. +# TIMEOUT env variable is used by nvcr.io/nvidia/pytorch:xx for another purpose +# rendering TIMEOUT defined by uvicorn impossible to use correctly +# We're overriding it to be renamed UVICORN_TIMEOUT +# UVICORN_TIMEOUT is a useful variable for very large models that take more +# than 30s (the default) to load in memory. +# If UVICORN_TIMEOUT is too low, uvicorn will simply never loads as it will +# kill workers all the time before they finish. +RUN sed -i 's/TIMEOUT/UVICORN_TIMEOUT/g' /gunicorn_conf.py +COPY ./app /app/app diff --git a/docker_images/thirdai/app/__init__.py b/docker_images/thirdai/app/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/docker_images/thirdai/app/main.py b/docker_images/thirdai/app/main.py new file mode 100644 index 00000000..ae342596 --- /dev/null +++ b/docker_images/thirdai/app/main.py @@ -0,0 +1,92 @@ +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, TextClassificationPipeline, TokenClassificationPipeline +from starlette.applications import Starlette +from starlette.middleware import Middleware +from starlette.middleware.gzip import GZipMiddleware +from starlette.routing import Route + + +TASK = os.getenv("TASK") +MODEL_ID = os.getenv("MODEL_ID") + + +logger = logging.getLogger(__name__) + + +# Add the allowed tasks +# Supported tasks are: +# - text-generation +# - text-classification +# - token-classification +# - translation +# - summarization +# - automatic-speech-recognition +# - ... +# For instance +# from app.pipelines import AutomaticSpeechRecognitionPipeline +# ALLOWED_TASKS = {"automatic-speech-recognition": AutomaticSpeechRecognitionPipeline} +# You can check the requirements and expectations of each pipelines in their respective +# directories. Implement directly within the directories. +ALLOWED_TASKS: Dict[str, Type[Pipeline]] = { + "text-classification": TextClassificationPipeline, + "token-classification": TokenClassificationPipeline, +} + + +@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) + + +routes = [ + Route("/{whatever:path}", status_ok), + Route("/{whatever:path}", pipeline_route, methods=["POST"]), +] + +middleware = [Middleware(GZipMiddleware, minimum_size=1000)] +if os.environ.get("DEBUG", "") == "1": + from starlette.middleware.cors import CORSMiddleware + + middleware.append( + Middleware( + CORSMiddleware, + allow_origins=["*"], + allow_headers=["*"], + allow_methods=["*"], + ) + ) + +app = Starlette(routes=routes, middleware=middleware) + + +@app.on_event("startup") +async def startup_event(): + logger = logging.getLogger("uvicorn.access") + handler = logging.StreamHandler() + handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")) + logger.handlers = [handler] + + # Link between `api-inference-community` and framework code. + app.get_pipeline = get_pipeline + try: + get_pipeline() + except Exception: + # We can fail so we can show exception later. + pass + + +if __name__ == "__main__": + try: + get_pipeline() + except Exception: + # We can fail so we can show exception later. + pass diff --git a/docker_images/thirdai/app/pipelines/__init__.py b/docker_images/thirdai/app/pipelines/__init__.py new file mode 100644 index 00000000..150e3a9e --- /dev/null +++ b/docker_images/thirdai/app/pipelines/__init__.py @@ -0,0 +1,4 @@ +from app.pipelines.base import Pipeline, PipelineException + +from app.pipelines.text_classification import TextClassificationPipeline +from app.pipelines.token_classification import TokenClassificationPipeline diff --git a/docker_images/thirdai/app/pipelines/base.py b/docker_images/thirdai/app/pipelines/base.py new file mode 100644 index 00000000..74a9080f --- /dev/null +++ b/docker_images/thirdai/app/pipelines/base.py @@ -0,0 +1,16 @@ +from abc import ABC, abstractmethod +from typing import Any + + +class Pipeline(ABC): + @abstractmethod + def __init__(self, model_id: str): + raise NotImplementedError("Pipelines should implement an __init__ method") + + @abstractmethod + def __call__(self, inputs: Any) -> Any: + raise NotImplementedError("Pipelines should implement a __call__ method") + + +class PipelineException(Exception): + pass \ No newline at end of file diff --git a/docker_images/thirdai/app/pipelines/text_classification.py b/docker_images/thirdai/app/pipelines/text_classification.py new file mode 100644 index 00000000..f0cc89d1 --- /dev/null +++ b/docker_images/thirdai/app/pipelines/text_classification.py @@ -0,0 +1,37 @@ +from typing import Dict, List + +from app.pipelines import Pipeline +from thirdai import bolt +import numpy as np + +from huggingface_hub import hf_hub_download + + +class TextClassificationPipeline(Pipeline): + def __init__( + self, + model_id: str, + ): + model_path = hf_hub_download(model_id, "model.bin", library_name="thirdai") + self.model = bolt.UniversalDeepTransformer.load(model_path) + + def __call__(self, inputs: str) -> List[Dict[str, float]]: + """ + Args: + inputs (:obj:`str`): + a string containing some text + Return: + A :obj:`list`:. The object returned should be a list of one list like [[{"label": 0.9939950108528137}]] containing: + - "label": A string representing what the label/class is. There can be multiple labels. + - "score": A score between 0 and 1 describing how confident the model is for this label/class. + """ + outputs = self.model.predict({self.model.text_dataset_config().text_column: inputs}) + + if len(outputs) == 0: + return [] + + if isinstance(outputs[0], tuple): + return [[{str(outputs[0][0]): outputs[0][1]}]] + else: + index = np.argmax(outputs) + return [[{self.model.class_name(index): outputs[index]}]] diff --git a/docker_images/thirdai/app/pipelines/token_classification.py b/docker_images/thirdai/app/pipelines/token_classification.py new file mode 100644 index 00000000..c14e9505 --- /dev/null +++ b/docker_images/thirdai/app/pipelines/token_classification.py @@ -0,0 +1,52 @@ +from typing import Any, Dict, List + +from app.pipelines import Pipeline +from thirdai import bolt + +from huggingface_hub import hf_hub_download + + +class TokenClassificationPipeline(Pipeline): + def __init__( + self, + model_id: str, + ): + print("LMFAO", model_id) + model_path = hf_hub_download(model_id, "model.bin", library_name="thirdai") + self.model = bolt.UniversalDeepTransformer.NER.load(model_path) + + def __call__(self, inputs: str) -> List[Dict[str, Any]]: + """ + Args: + inputs (:obj:`str`): + a string containing some text + Return: + A :obj:`list`:. The object returned should be like [{"entity_group": "XXX", "word": "some word", "start": 3, "end": 6, "score": 0.82}] containing : + - "entity_group": A string representing what the entity is. + - "word": A rubstring of the original string that was detected as an entity. + - "start": the offset within `input` leading to `answer`. context[start:stop] == word + - "end": the ending offset within `input` leading to `answer`. context[start:stop] === word + - "score": A score between 0 and 1 describing how confident the model is for this entity. + """ + split_inputs = inputs.split(" ") + + outputs = self.model.predict(split_inputs) + + entities = [] + offset = 0 + for entity_results, word in zip(outputs, split_inputs): + best_prediction = entity_results[0] + + current_entity = { + "entity_group": best_prediction[0], + "word": word, + "start": offset, + "end": offset + len(word), + "score": best_prediction[1], + } + + entities.append(current_entity) + + offset += len(word) + 1 + + return entities \ No newline at end of file diff --git a/docker_images/thirdai/prestart.sh b/docker_images/thirdai/prestart.sh new file mode 100644 index 00000000..fdcc34b4 --- /dev/null +++ b/docker_images/thirdai/prestart.sh @@ -0,0 +1 @@ +python app/main.py diff --git a/docker_images/thirdai/requirements.txt b/docker_images/thirdai/requirements.txt new file mode 100644 index 00000000..8709d83b --- /dev/null +++ b/docker_images/thirdai/requirements.txt @@ -0,0 +1,5 @@ +starlette==0.27.0 +api-inference-community==0.0.32 +huggingface_hub==0.11.0 +thirdai==0.8.5 +numpy==1.25.2 \ No newline at end of file diff --git a/docker_images/thirdai/tests/__init__.py b/docker_images/thirdai/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/docker_images/thirdai/tests/samples/malformed.flac b/docker_images/thirdai/tests/samples/malformed.flac new file mode 100644 index 00000000..06d74050 Binary files /dev/null and b/docker_images/thirdai/tests/samples/malformed.flac differ diff --git a/docker_images/thirdai/tests/samples/plane.jpg b/docker_images/thirdai/tests/samples/plane.jpg new file mode 100644 index 00000000..9e0cc585 Binary files /dev/null and b/docker_images/thirdai/tests/samples/plane.jpg differ diff --git a/docker_images/thirdai/tests/samples/plane2.jpg b/docker_images/thirdai/tests/samples/plane2.jpg new file mode 100644 index 00000000..b1f076e3 Binary files /dev/null and b/docker_images/thirdai/tests/samples/plane2.jpg differ diff --git a/docker_images/thirdai/tests/samples/sample1.flac b/docker_images/thirdai/tests/samples/sample1.flac new file mode 100644 index 00000000..837cd984 Binary files /dev/null and b/docker_images/thirdai/tests/samples/sample1.flac differ diff --git a/docker_images/thirdai/tests/samples/sample1.webm b/docker_images/thirdai/tests/samples/sample1.webm new file mode 100644 index 00000000..7c449a27 Binary files /dev/null and b/docker_images/thirdai/tests/samples/sample1.webm differ diff --git a/docker_images/thirdai/tests/samples/sample1_dual.ogg b/docker_images/thirdai/tests/samples/sample1_dual.ogg new file mode 100644 index 00000000..1e8524e0 Binary files /dev/null and b/docker_images/thirdai/tests/samples/sample1_dual.ogg differ diff --git a/docker_images/thirdai/tests/test_api.py b/docker_images/thirdai/tests/test_api.py new file mode 100644 index 00000000..3b58a98e --- /dev/null +++ b/docker_images/thirdai/tests/test_api.py @@ -0,0 +1,60 @@ +import os +from typing import Dict +from unittest import TestCase, skipIf + +from app.main import ALLOWED_TASKS, get_pipeline + + +# Must contain at least one example of each implemented pipeline +# Tests do not check the actual values of the model output, so small dummy +# models are recommended for faster tests. +TESTABLE_MODELS: Dict[str, str] = { + "text-classification": "thirdai/Classification", + "token-classification": "thirdai/NamedEntityRecognition", +} + + +ALL_TASKS = { + "audio-classification", + "audio-to-audio", + "automatic-speech-recognition", + "feature-extraction", + "image-classification", + "question-answering", + "sentence-similarity", + "speech-segmentation", + "tabular-classification", + "tabular-regression", + "text-to-image", + "text-to-speech", + "token-classification", + "conversational", + "feature-extraction", + "sentence-similarity", + "fill-mask", + "table-question-answering", + "summarization", + "text2text-generation", + "text-classification", + "zero-shot-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/docker_images/thirdai/tests/test_api_text_classification.py b/docker_images/thirdai/tests/test_api_text_classification.py new file mode 100644 index 00000000..abaa5222 --- /dev/null +++ b/docker_images/thirdai/tests/test_api_text_classification.py @@ -0,0 +1,86 @@ +import json +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( + "text-classification" not in ALLOWED_TASKS, + "text-classification not implemented", +) +class TextClassificationTestCase(TestCase): + def setUp(self): + model_id = TESTABLE_MODELS["text-classification"] + self.old_model_id = os.getenv("MODEL_ID") + self.old_task = os.getenv("TASK") + os.environ["MODEL_ID"] = model_id + os.environ["TASK"] = "text-classification" + from app.main import app + + 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/docker_images/thirdai/tests/test_api_token_classification.py b/docker_images/thirdai/tests/test_api_token_classification.py new file mode 100644 index 00000000..8f5f8729 --- /dev/null +++ b/docker_images/thirdai/tests/test_api_token_classification.py @@ -0,0 +1,83 @@ +import json +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( + "token-classification" not in ALLOWED_TASKS, + "token-classification not implemented", +) +class TokenClassificationTestCase(TestCase): + def setUp(self): + model_id = TESTABLE_MODELS["token-classification"] + self.old_model_id = os.getenv("MODEL_ID") + self.old_task = os.getenv("TASK") + os.environ["MODEL_ID"] = model_id + os.environ["TASK"] = "token-classification" + from app.main import app + + 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 = "Hello, my name is John and I live in New York" + + 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( + set(k for el in content for k in el.keys()), + {"entity_group", "word", "start", "end", "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( + set(k for el in content for k in el.keys()), + {"entity_group", "word", "start", "end", "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/docker_images/thirdai/tests/test_docker_build.py b/docker_images/thirdai/tests/test_docker_build.py new file mode 100644 index 00000000..ca5dd65d --- /dev/null +++ b/docker_images/thirdai/tests/test_docker_build.py @@ -0,0 +1,23 @@ +import os +import subprocess +from unittest import TestCase + + +class cd: + """Context manager for changing the current working directory""" + + def __init__(self, newPath): + self.newPath = os.path.expanduser(newPath) + + def __enter__(self): + self.savedPath = os.getcwd() + os.chdir(self.newPath) + + def __exit__(self, etype, value, traceback): + os.chdir(self.savedPath) + + +class DockerBuildTestCase(TestCase): + def test_can_build_docker_image(self): + with cd(os.path.dirname(os.path.dirname(__file__))): + subprocess.check_output(["docker", "build", "."]) diff --git a/tests/test_dockers.py b/tests/test_dockers.py index 7ad2dec6..6fa7870c 100644 --- a/tests/test_dockers.py +++ b/tests/test_dockers.py @@ -315,6 +315,18 @@ def test_timm(self): self.framework_docker_test("timm", "image-classification", "sgugger/resnet50d") self.framework_invalid_test("timm") + def test_thirdai(self): + self.framework_docker_test( + "thirdai", + "text-classification", + "thirdai/Classification", + ) + self.framework_docker_test( + "thirdai", + "token-classification", + "thirdai/NamedEntityRecognition", + ) + def test_diffusers(self): self.framework_docker_test( "diffusers",