feat: inference API cleanup (kaito-project#233)

This PR does some cleanup on the inference API. Introduces: 1. ModelConfig Dataclass 2. Allow additional user args 3. GenerateKwargs Model 4. More descriptive comments/metadata for model fields 5. Rename inference-api.py -> inference_api.py - needed for naming convention/importability 6. Introduce tests for inference_api.py This PR requires bumping up images and having Kaito use updated image tag (in this case 0.0.2).
bangqipropel · Feb 21, 2024 · 68c972d · 68c972d
1 parent 7f50583
commit 68c972d
Show file tree

Hide file tree

Showing 23 changed files with 497 additions and 214 deletions.
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -38,6 +38,10 @@ jobs:
         run: |
           make unit-test
 
+      - name: Run inference api unit tests
+        run: |
+          make inference-api-e2e
+
       - name: Upload Codecov report
         uses: codecov/codecov-action@v4
         with:

diff --git a/Makefile b/Makefile
@@ -83,6 +83,10 @@ unit-test: ## Run unit tests.
 	go test -v $(shell go list ./pkg/... ./api/... | grep -v /vendor) -race -coverprofile=coverage.txt -covermode=atomic
 	go tool cover -func=coverage.txt
 
+inference-api-e2e: 
+	pip install -r presets/inference/text-generation/requirements.txt
+	pytest -o log_cli=true -o log_cli_level=INFO .
+
 $(E2E_TEST):
 	(cd test/e2e && go test -c . -o $(E2E_TEST))
 

diff --git a/docker/presets/tfs-onnx/Dockerfile b/docker/presets/tfs-onnx/Dockerfile
@@ -17,7 +17,7 @@ RUN echo $VERSION > /workspace/tfs/version.txt
 COPY kaito/presets/inference/${MODEL_TYPE}/requirements.txt /workspace/tfs/requirements.txt
 RUN pip install --no-cache-dir -r requirements.txt
 
-COPY kaito/presets/inference/${MODEL_TYPE}/inference-api.py /workspace/tfs/inference-api.py
+COPY kaito/presets/inference/${MODEL_TYPE}/inference_api.py /workspace/tfs/inference_api.py
 
 # Convert to ONNX Runtime
 # RUN python convert_to_onnx.py ${MODEL_NAME} 

diff --git a/docker/presets/tfs/Dockerfile b/docker/presets/tfs/Dockerfile
@@ -1,4 +1,4 @@
-FROM python:3.8-slim
+FROM python:3.10-slim
 
 ARG WEIGHTS_PATH
 ARG MODEL_TYPE
@@ -16,7 +16,7 @@ RUN echo $VERSION > /workspace/tfs/version.txt
 COPY kaito/presets/inference/${MODEL_TYPE}/requirements.txt /workspace/tfs/requirements.txt
 RUN pip install --no-cache-dir -r requirements.txt
 
-COPY kaito/presets/inference/${MODEL_TYPE}/inference-api.py /workspace/tfs/inference-api.py
+COPY kaito/presets/inference/${MODEL_TYPE}/inference_api.py /workspace/tfs/inference_api.py
 
 # Copy the entire model weights to the weights directory
 COPY ${WEIGHTS_PATH} /workspace/tfs/weights
diff --git a/...ts/inference/llama2-chat/inference-api.py → ...ts/inference/llama2-chat/inference_api.py b/...ts/inference/llama2-chat/inference-api.py → ...ts/inference/llama2-chat/inference_api.py
diff --git a/...erence/llama2-completion/inference-api.py → ...erence/llama2-completion/inference_api.py b/...erence/llama2-completion/inference-api.py → ...erence/llama2-completion/inference_api.py
diff --git a/presets/inference/text-generation/inference-api.py b/presets/inference/text-generation/inference-api.py