From 39a334adf84952e72d46c54a5c580105014b15e4 Mon Sep 17 00:00:00 2001 From: bjwswang Date: Fri, 1 Mar 2024 08:54:09 +0000 Subject: [PATCH] feat(worker): add a new core-library-cli runner to host reranking models Signed-off-by: bjwswang --- .github/workflows/image_build.yml | 61 ++++++++++++++++++++++++ .github/workflows/image_build_test.yml | 42 +++++++++++++++++ docker/Dockerfile | 29 ++++++++++++ libs/cli/README.md | 3 +- libs/cli/kubeagi_cli/cli.py | 1 - libs/cli/kubeagi_cli/serve/reranking.py | 62 ++++++++++++++++++++++--- libs/cli/kubeagi_cli/serve/whisper.py | 20 ++++++++ libs/cli/kubeagi_cli/server.py | 15 +++--- libs/cli/pyproject.toml | 7 ++- 9 files changed, 223 insertions(+), 17 deletions(-) create mode 100644 .github/workflows/image_build.yml create mode 100644 .github/workflows/image_build_test.yml create mode 100644 docker/Dockerfile create mode 100644 libs/cli/kubeagi_cli/serve/whisper.py diff --git a/.github/workflows/image_build.yml b/.github/workflows/image_build.yml new file mode 100644 index 0000000..41a4710 --- /dev/null +++ b/.github/workflows/image_build.yml @@ -0,0 +1,61 @@ +name: Build KubeAGI Core Library CLI Image + +on: + push: + branches: [main] + paths: + - 'libs/**' + workflow_dispatch: +env: + PYTHON_INDEX_URL: https://pypi.org/simple + +jobs: + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - name: Set Variable + id: set-env + run: | + TAG=$(git describe --tags --abbrev=0 --match 'v*' 2> /dev/null) || true + if [ -z "$TAG" ]; then + echo "No tag found, use v0.0.1 as default" + TAG=v0.0.1 + fi + echo "TAG=${TAG}" >> $GITHUB_OUTPUT + echo "DATE=$(TZ=Asia/Shanghai date +'%Y%m%d')" >> $GITHUB_OUTPUT + - name: Show Variable + run: echo "varibables ${{ steps.set-env.outputs.TAG }}-${{ steps.set-env.outputs.DATE }}" + - name: Set up QEMU + uses: docker/setup-qemu-action@v3 + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + with: + buildkitd-flags: --debug + config-inline: | + [worker.oci] + max-parallelism = 1 + - name: Login to the dockerhub Registry + uses: docker/login-action@v3 + with: + username: ${{ secrets.DOCKER_USERNAME }} + password: ${{ secrets.DOCKER_TOKEN }} + - uses: benjlevesque/short-sha@v2.2 + name: Get short commit sha + id: short-sha + - name: Build and push Evaluation Image + id: push-eval + uses: docker/build-push-action@v5 + with: + context: . + file: docker/Dockerfile + platforms: linux/amd64,linux/arm64 + tags: | + kubeagi/core-library-cli:latest + kubeagi/core-library-cli:${{ steps.set-env.outputs.TAG }} + kubeagi/core-library-cli:${{ steps.set-env.outputs.TAG }}-${{ steps.set-env.outputs.DATE }}-${{ steps.short-sha.outputs.sha }} + push: true + build-args: | + PYTHON_INDEX_URL=${{ env.PYTHON_INDEX_URL }} diff --git a/.github/workflows/image_build_test.yml b/.github/workflows/image_build_test.yml new file mode 100644 index 0000000..d8e80df --- /dev/null +++ b/.github/workflows/image_build_test.yml @@ -0,0 +1,42 @@ +name: Test Build KubeAGI Core Library CLI Image + +on: + pull_request: + branches: [main] + paths: + - 'libs/**' + workflow_dispatch: +env: + PYTHON_INDEX_URL: https://pypi.org/simple + +jobs: + test_image_build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - name: Set up QEMU + uses: docker/setup-qemu-action@v3 + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@v3 + with: + buildkitd-flags: --debug + config-inline: | + [worker.oci] + max-parallelism = 1 + - name: Set up GCC + uses: egor-tensin/setup-gcc@v1 + with: + version: latest + platform: x64 + - name: Build core library cli image + id: push-worker + uses: docker/build-push-action@v5 + with: + context: . + file: docker/Dockerfile + platforms: linux/amd64,linux/arm64 + push: false + build-args: | + PYTHON_INDEX_URL=${{ env.PYTHON_INDEX_URL }} \ No newline at end of file diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 0000000..a3121e4 --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,29 @@ +ARG PY_VER=3.11 + +# python environment +FROM python:${PY_VER}-slim AS runner +ARG PACKAGE_REGISTRY="mirrors.tuna.tsinghua.edu.cn" +RUN sed -i 's/deb.debian.org/mirrors.tuna.tsinghua.edu.cn/g' /etc/apt/sources.list.d/debian.sources + +ENV TZ=Asia/Shanghai +RUN export DEBIAN_FRONTEND=noninteractive \ + && apt-get update \ + && apt-get install -y tzdata \ + && ln -fs /usr/share/zoneinfo/Asia/Shanghai /etc/localtime \ + && dpkg-reconfigure --frontend noninteractive tzdata + +RUN apt-get install gcc python3-dev + +# Official: https://pypi.org/simple +ARG PYTHON_INDEX_URL=https://pypi.mirrors.ustc.edu.cn/simple/ +COPY libs /libs +RUN python -m pip install ragas langchain==0.0.1 -i ${PYTHON_INDEX_URL} +WORKDIR /libs/core +RUN pip install -e . -i ${PYTHON_INDEX_URL} + +WORKDIR /libs/cli +RUN pip install -e . -i ${PYTHON_INDEX_URL} + +ENV RERANKING_MODEL_PATH=BAAI/bge-reranker-large + +CMD [ "python","kubeagi_cli/cli.py","serve" ] diff --git a/libs/cli/README.md b/libs/cli/README.md index 153c48e..32baca9 100644 --- a/libs/cli/README.md +++ b/libs/cli/README.md @@ -51,4 +51,5 @@ kubeagi-cli evaluate context_precision ~/core-library/examples/testdata/ragas.cs ``` The above command will run the rag evaluation with metrics `context_precision` and test dataset `~/core-library/examples/testdata/ragas.csv` with the help -of OpenAI. \ No newline at end of file +of OpenAI. + diff --git a/libs/cli/kubeagi_cli/cli.py b/libs/cli/kubeagi_cli/cli.py index 133a23d..a7b7a04 100644 --- a/libs/cli/kubeagi_cli/cli.py +++ b/libs/cli/kubeagi_cli/cli.py @@ -15,7 +15,6 @@ import os import typer from typing_extensions import Annotated -from typing import Optional from kubeagi_core.evaluation.ragas_eval import RagasEval diff --git a/libs/cli/kubeagi_cli/serve/reranking.py b/libs/cli/kubeagi_cli/serve/reranking.py index 21d344c..d1568cb 100644 --- a/libs/cli/kubeagi_cli/serve/reranking.py +++ b/libs/cli/kubeagi_cli/serve/reranking.py @@ -15,28 +15,76 @@ from typing import List from FlagEmbedding import FlagReranker +from BCEmbedding import RerankerModel +from abc import ABC, abstractmethod -class Reranking: +class BaseReranking(ABC): """ - The Reranking is used to run reranking models like bge-reranker-large(https://huggingface.co/BAAI/bge-reranker-large) + The Reranking is used to run reranking models. + """ + + @abstractmethod + def run(self, pairs: List[List[str]]): + """run reranking models to rerank pairs.""" + + +class BGEReranking(BaseReranking): + """ + The BGEReranking is used to run reranking models with FlagEmbedding like sbge-reranker-large(https://huggingface.co/BAAI/bge-reranker-large) """ - _model_path: str _reranker: FlagReranker def __init__( self, - model_path: str, + model_name_or_path: str, + ): + self._reranker = FlagReranker( + model_name_or_path=model_name_or_path, use_fp16=False + ) + + # run bge reranking model + def run(self, pairs: List[List[str]]): + if len(pairs) > 0: + result = self._reranker.compute_score(pairs) + if isinstance(result, float): + result = [result] + return result + else: + return None + + +class BCEReranking(BaseReranking): + """ + The BGEReranking is used to run reranking models with BCEEmbedding from https://github.com/netease-youdao/BCEmbedding + """ + + _reranker: RerankerModel + + def __init__( + self, + model_name_or_path: str, ): - self._model_path = model_path + self._reranker = RerankerModel( + model_name_or_path=model_name_or_path, use_fp16=False + ) + # run bge reranking model def run(self, pairs: List[List[str]]): - reranker = FlagReranker(self._model_path, use_fp16=False) if len(pairs) > 0: - result = reranker.compute_score(pairs) + result = self._reranker.compute_score(pairs) if isinstance(result, float): result = [result] return result else: return None + + +def select_reranking(model_name_or_path: str) -> BaseReranking: + if "bge" in model_name_or_path.lower(): + return BGEReranking(model_name_or_path) + if "bce" in model_name_or_path.lower(): + return BCEReranking(model_name_or_path) + + raise ValueError(f"No valid reranking runner for {model_name_or_path}") diff --git a/libs/cli/kubeagi_cli/serve/whisper.py b/libs/cli/kubeagi_cli/serve/whisper.py new file mode 100644 index 0000000..8e89869 --- /dev/null +++ b/libs/cli/kubeagi_cli/serve/whisper.py @@ -0,0 +1,20 @@ +# Copyright 2024 KubeAGI. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +# TODO +class Whisper: + """ + The Whisper is used to run OpenAI's whisper model + """ diff --git a/libs/cli/kubeagi_cli/server.py b/libs/cli/kubeagi_cli/server.py index c603ab1..41187d2 100644 --- a/libs/cli/kubeagi_cli/server.py +++ b/libs/cli/kubeagi_cli/server.py @@ -16,8 +16,8 @@ from typing import Optional, List from fastapi import FastAPI -from pydantic import BaseModel -from serve.reranking import Reranking +from pydantic import BaseModel, Field +from serve.reranking import select_reranking app = FastAPI() @@ -28,14 +28,17 @@ def health(): class RerankingInput(BaseModel): + model_name_or_path: Optional[str] = Field(default=os.getenv("RERANKING_MODEL_PATH")) question: str answers: Optional[List[str]] @app.post("/api/v1/reranking") -def reranking(input_docs: RerankingInput): +def reranking(input: RerankingInput): + # select reranking models based on model path + reranker = select_reranking(input.model_name_or_path) + # prepare reranking pairs pairs = [] - for answer in input_docs.answers: - pairs.append([input_docs.question, answer]) - reranker = Reranking(model_path=os.getenv("RERANKING_MODEL_PATH")) + for answer in input.answers: + pairs.append([input.question, answer]) return reranker.run(pairs) diff --git a/libs/cli/pyproject.toml b/libs/cli/pyproject.toml index 539799d..c14020c 100644 --- a/libs/cli/pyproject.toml +++ b/libs/cli/pyproject.toml @@ -17,8 +17,11 @@ classifiers = [ ] dependencies = [ "typer==0.9.0", - "FlagEmbedding=1.2.5", - "kubeagi-core=0.0.1", + "fastapi==0.109.0", + "uvicorn==0.27.0", + "FlagEmbedding==1.2.3", + "BCEmbedding==0.1.3", + "kubeagi-core==0.0.1", ] [project.optional-dependencies]