support for embedding api (OpenAgentsInc#22)

Author: earonesty

.github/workflows/main.yml

@@ -0,0 +1,50 @@
+name: Python package
+
+on: [push]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+env:
+  BILLING_URL: ${{ vars.BILLING_URL }}
+  SECRET_KEY: ${{ vars.SECRET_KEY }}
+  BYPASS_TOKEN: ${{ vars.BYPASS_TOKEN }}
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.11"]
+    env:
+      COVERAGE_FILE: ".coverage.${{ matrix.info.os }}.${{ matrix.info.python }}.${{ matrix.info.mysql }}"
+    steps:
+      - uses: actions/checkout@v3
+      - name: Install poetry
+        run: pipx install poetry
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: 'poetry'
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install ruff pytest poetry coverage
+          poetry lock --check
+          poetry install --with=dev-onnx
+      - name: Lint with ruff
+        run: |
+          # stop the build if there are Python syntax errors or undefined names
+          ruff --output-format=github --target-version=py311 .
+      - name: Test with pytest
+        run: |
+          ./run_tests.sh
+      - uses: actions/upload-artifact@v3
+        with:
+          name: coverage
+          path: |
+            ${{ env.COVERAGE_FILE }}
+            htmlcov
+

.gitignore

@@ -2,10 +2,11 @@ ai_worker/version.py
 
 # lib stuff
 CLBlast/
-build-cuda/
-build-opencl/
+build-*/
 bin/
 *.gguf
+local_cache/
+tests/local_cache/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

ai_worker/fast_embed.py

@@ -0,0 +1,70 @@
+import json
+import os
+import logging as log
+from typing import Optional
+
+DEFAULT_MAX_LENGTH = 512
+DEFAULT_MODEL = "BAAI/bge-base-en-v1.5"
+MODEL_PREFIX = "fastembed:"
+
+
+class _FastEmbed:
+    def __init__(self, cls, conf):
+        self.conf = conf
+        self.embedding_class = cls
+        self.embedding_model = cls(model_name=DEFAULT_MODEL, max_length=DEFAULT_MAX_LENGTH)
+
+    def embed(self, req: dict):
+        model = req["model"]
+
+        if model.startswith(MODEL_PREFIX):
+            model = model[len(MODEL_PREFIX):]
+
+        max_length = req.get("max_length", 512)
+
+        if self.embedding_model.model_name != model or self.embedding_model._max_length != max_length:  # noqa
+            # swap out model
+            self.embedding_model = self.embedding_class(model_name=model, max_length=max_length)
+
+        docs = req["input"]
+        if isinstance(docs, str):
+            docs = [docs]
+
+        # todo: better toks count
+        toks = int(len(json.dumps(docs)) / 2.5)
+
+        res = {
+            "object": "list",
+            "model": model,
+            "usage": {
+                "prompt_tokens": toks,
+                "total_tokens": toks
+            }
+        }
+
+        embed = [
+            dict(
+                object="embedding",
+                embedding=nda.tolist(),
+                index=i
+            )
+            for i, nda in enumerate(self.embedding_model.embed(docs, parallel=0))
+        ]
+
+        res["data"] = embed
+
+        return res
+
+
+def FastEmbed(*a) -> Optional[_FastEmbed]:
+    try:
+        from fastembed.embedding import FlagEmbedding as Embedding
+        import onnxruntime as ort
+        if ort.get_device() != "GPU" and not os.environ.get("CI"):
+            log.warning("fast embed not enabled, ort runtime does not see the GPU")
+            return None
+        return _FastEmbed(Embedding, *a)
+    except ImportError:
+        if os.environ.get("GPUTOPIA_DEBUG_IMPORT"):
+            log.exception("fast embed not enabled")
+        return None

ai_worker/fine_tune.py

@@ -1,12 +1,10 @@
 import gc
-import hashlib
 import json
 import asyncio
 import threading
 import logging
 import os
 import random
-import tarfile
 import shutil
 
 import transformers
@@ -17,7 +15,7 @@
 from peft import prepare_model_for_kbit_training, PeftModel, LoraConfig, get_peft_model
 from accelerate import FullyShardedDataParallelPlugin, Accelerator
 from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig
-from ai_worker.util import quantize_gguf, url_to_tempfile, user_ft_name_to_url
+from ai_worker.util import quantize_gguf, url_to_tempfile, user_ft_name_to_url, gzip
 
 from ai_worker.util import b64enc
 from gguf_loader.convert import main as gguf_main
@@ -27,14 +25,6 @@
 log = logging.getLogger(__name__)
 
 
-def gzip(folder):
-    """tar gz the folder to 'folder.tar.gz', removes the folder"""
-    base_folder_name = os.path.basename(folder)
-    with tarfile.open(f"{folder}.tar.gz", 'w:gz') as archive:
-        archive.add(folder, arcname=base_folder_name)
-    return f"{folder}.tar.gz"
-
-
 class FineTuner:
     def __init__(self, conf):
         self.conf = conf

ai_worker/gguf_reader.py

@@ -131,8 +131,8 @@ def read_header(self):
 
         self.gguf_version = struct.unpack("<I", self.fin.read(4))[0]
 
-        if self.gguf_version != 2:
-            raise ValueError("Can only summarize version 2 files")
+        if self.gguf_version < 2:
+            raise ValueError("Can only summarize version 2/3 files, got version %s" % self.gguf_version)
 
         self.ti_data_count = struct.unpack("<Q", self.fin.read(8))[0]
         self.kv_data_count = struct.unpack("<Q", self.fin.read(8))[0]

ai_worker/jsonlines.py

@@ -32,11 +32,13 @@
 
 try:
     from .fine_tune import FineTuner
-except ImportError as ex:
+except ImportError:
     if os.environ.get("GPUTOPIA_DEBUG_IMPORT"):
         log.exception("fine tuning not enabled")
     FineTuner = None
 
+from .fast_embed import FastEmbed, MODEL_PREFIX
+
 from gguf_loader.main import get_size
 
 from .gguf_reader import GGUFReader
@@ -134,11 +136,14 @@ def __init__(self, conf: Config):
         self.llama = None
         self.llama_model = None
         self.llama_cli: Optional[AsyncClient] = None
+        
         if FineTuner:
             self.fine_tuner = FineTuner(self.conf)
         else:
             self.fine_tuner = None
 
+        self.fast_embed = FastEmbed(self.conf)
+
     def _gen_or_load_priv(self) -> None:
         if not self.conf.privkey:
             cfg = self.conf.config
@@ -215,6 +220,10 @@ async def guess_layers(self, model_path):
 
         for gpu in info.nv_gpus:
             tot_mem += gpu.memory * 1000000
+        
+        if tot_mem == 0:
+            for gpu in info.cl_gpus:
+                tot_mem += gpu.memory * 1000000
 
         if est_ram > tot_mem:
             est_layers = tot_mem // (est_ram / layers)
@@ -227,9 +236,12 @@ async def guess_layers(self, model_path):
         return max(0, est_layers - self.conf.layer_offset)
 
     async def load_model(self, name):
+        assert name, "No model name"
         if name == self.llama_model:
             return
+        
         log.debug("loading model: %s", name)
+        
         model_path = await self.get_model(name)
 
         if llama_cpp.server.app.llama:
@@ -243,6 +255,7 @@ async def load_model(self, name):
                                  embedding=True, cache=True, port=8181,
                                  main_gpu=self.conf.main_gpu, tensor_split=sp)
         self.llama = create_llama_app(settings)
+        assert self.llama, "Load llama failed.   Try lowering layers."
         self.llama_cli = AsyncClient(app=self.llama, base_url="http://test")
         self.llama_model = name
 
@@ -256,6 +269,9 @@ def _get_connect_info(self) -> ConnectMessage:
         if self.fine_tuner:
             caps += ["llama-fine-tune"]
 
+        if self.fast_embed:
+            caps += ["fast-embed"]
+
         connect_msg = ConnectMessage(
             worker_version=VERSION,
             capabilities=caps,
@@ -366,6 +382,9 @@ async def run_one(self):
                 async for event in self.fine_tuner.fine_tune(req.openai_req):
                     await self.ws_send(json.dumps(event), True)
                 await self.ws_send("{}")
+            elif req.openai_url == "/v1/embeddings" and model.startswith(MODEL_PREFIX):
+                res = self.fast_embed.embed(req.openai_req)
+                await self.ws_send(json.dumps(res), True)
             elif req.openai_req.get("stream"):
                 await self.load_model(model)
                 async with aconnect_sse(self.llama_cli, "POST", req.openai_url, json=req.openai_req) as sse:

ai_worker/main.py

@@ -2,8 +2,6 @@
 
 from ai_worker.util import quantize_gguf, GGML_INVERSE_MAP
 
-import argparse
-
 # List of quantization levels
 quantization_levels = list(GGML_INVERSE_MAP.keys())
 

ai_worker/quantize_main.py

@@ -1,6 +1,7 @@
 import base64
 import hashlib
 import os
+import tarfile
 
 import llama_cpp
 
@@ -61,3 +62,11 @@ def url_to_tempfile(conf, url):
     name = hashlib.md5(url.encode()).hexdigest()
     output_file = os.path.join(conf.tmp_dir, name)
     return output_file
+
+
+def gzip(folder):
+    """tar gz the folder to 'folder.tar.gz', removes the folder"""
+    base_folder_name = os.path.basename(folder)
+    with tarfile.open(f"{folder}.tar.gz", 'w:gz') as archive:
+        archive.add(folder, arcname=base_folder_name)
+    return f"{folder}.tar.gz"

ai_worker/util.py

@@ -1 +1 @@
-VERSION = '0.2.3'
+VERSION = '0.3.0'

ai_worker/version.py

@@ -16,6 +16,10 @@ if [ "$gpu" == "cuda-torch" ]; then
     opts=""
 fi
 
+with_onnx=""
+if [ "$gpu" == "cuda" ]; then
+    with_torch="--with onnx"
+fi
 
 set -o xtrace
 
@@ -36,11 +40,4 @@ python build-version.py
 
 ./pyinstaller.sh $gpu-$arch $opts
 
-if [ "$gpu" == "cuda-torch" ]; then
-    pushd dist
-    tar cvf - gputopia-worker-$gpu-$arch/ | pigz -9 - > gputopia-worker-$gpu-$arch.tar.gz
-    rm -rf gputopia-worker-$gpu-$arch/
-    popd
-fi
-
 deactivate

build-bin.sh

@@ -3,11 +3,9 @@
 
 import argparse
 import concurrent.futures
-import copy
 import enum
 import faulthandler
 import functools
-import io
 import itertools
 import json
 import math
@@ -23,12 +21,11 @@
 from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
 from dataclasses import dataclass
 from pathlib import Path
-from typing import IO, TYPE_CHECKING, Any, Callable, Generator, Iterable, Literal, Sequence, TypeVar
+from typing import IO, TYPE_CHECKING, Any, Callable, Iterable, Literal, TypeVar
 
 import numpy as np
 from sentencepiece import SentencePieceProcessor  # type: ignore[import]
 
-import os
 from . import gguf
 
 if TYPE_CHECKING:
@@ -335,7 +332,6 @@ def __init__(self, fname_tokenizer: Path, fname_added_tokens: Path | None) -> No
 
     def bpe_tokens(self) -> Iterable[tuple[bytes, float, gguf.TokenType]]:
         tokenizer = self.bpe_tokenizer
-        from transformers.models.gpt2 import tokenization_gpt2  # type: ignore[import]
         reverse_vocab = {id: encoded_tok for encoded_tok, id in tokenizer.items()}
 
         for i, _ in enumerate(tokenizer):
@@ -850,7 +846,7 @@ def add_meta_vocab(self, vocab: Vocab) -> None:
         elif isinstance(vocab, BpeVocab):
             self.gguf.add_tokenizer_model("gpt2")
         else:
-            raise ValueError(f'Unknown vocab type: Not BpeVocab or SentencePieceVocab')
+            raise ValueError('Unknown vocab type: Not BpeVocab or SentencePieceVocab')
         self.gguf.add_token_list(tokens)
         self.gguf.add_token_scores(scores)
         self.gguf.add_token_types(toktypes)

build-mac.sh

@@ -2,7 +2,6 @@
 from __future__ import annotations
 
 import argparse
-import math
 import struct
 import sys
 from enum import IntEnum
@@ -425,7 +424,7 @@ def main():
     data = np.memmap(cfg.input, mode = 'r')
     model = GGMLModel()
     print('* Scanning GGML input file')
-    offset = model.load(data, 0)
+    model.load(data, 0)
     print(f'* GGML model hyperparameters: {model.hyperparameters}')
     vocab_override = None
     params_override = None