Integrate ThirdAI Library

.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}"'"}}'

.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

docker_images/thirdai/Dockerfile

@@ -0,0 +1,29 @@
+FROM --platform=linux/x86_64 tiangolo/uvicorn-gunicorn:python3.10
+LABEL maintainer="David Torres <[email protected]>"
+
+# 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

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

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

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

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]}]]

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

docker_images/thirdai/prestart.sh

@@ -0,0 +1 @@
+python app/main.py

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

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()