From c5be3f7acbb733acea49225f917154883307c710 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Fri, 10 Nov 2023 07:34:27 -0800 Subject: [PATCH] Make /v1/embeddings functional, add request/response types --- docs/12 - OpenAI API.md | 4 ++-- extensions/openai/cache_embedding_model.py | 4 ++-- extensions/openai/embeddings.py | 28 ++++++++++++---------- extensions/openai/models.py | 2 +- extensions/openai/script.py | 15 ++++++------ extensions/openai/typing.py | 13 ++++++++++ 6 files changed, 40 insertions(+), 26 deletions(-) diff --git a/docs/12 - OpenAI API.md b/docs/12 - OpenAI API.md index 1b0eccb5..05b4db02 100644 --- a/docs/12 - OpenAI API.md +++ b/docs/12 - OpenAI API.md @@ -211,7 +211,7 @@ The following environment variables can be used (they take precendence over ever | `OPENEDAI_KEY_PATH` | SSL key file path | key.pem | | `OPENEDAI_DEBUG` | Enable debugging (set to 1) | 1 | | `SD_WEBUI_URL` | WebUI URL (used by endpoint) | http://127.0.0.1:7861 | -| `OPENEDAI_EMBEDDING_MODEL` | Embedding model (if applicable) | all-mpnet-base-v2 | +| `OPENEDAI_EMBEDDING_MODEL` | Embedding model (if applicable) | sentence-transformers/all-mpnet-base-v2 | | `OPENEDAI_EMBEDDING_DEVICE` | Embedding device (if applicable) | cuda | #### Persistent settings with `settings.yaml` @@ -220,7 +220,7 @@ You can also set the following variables in your `settings.yaml` file: ``` openai-embedding_device: cuda -openai-embedding_model: all-mpnet-base-v2 +openai-embedding_model: "sentence-transformers/all-mpnet-base-v2" openai-sd_webui_url: http://127.0.0.1:7861 openai-debug: 1 ``` diff --git a/extensions/openai/cache_embedding_model.py b/extensions/openai/cache_embedding_model.py index 7f4f0806..8166cb4b 100644 --- a/extensions/openai/cache_embedding_model.py +++ b/extensions/openai/cache_embedding_model.py @@ -1,11 +1,11 @@ #!/usr/bin/env python3 # preload the embedding model, useful for Docker images to prevent re-download on config change # Dockerfile: -# ENV OPENEDAI_EMBEDDING_MODEL=all-mpnet-base-v2 # Optional +# ENV OPENEDAI_EMBEDDING_MODEL="sentence-transformers/all-mpnet-base-v2" # Optional # RUN python3 cache_embedded_model.py import os import sentence_transformers -st_model = os.environ.get("OPENEDAI_EMBEDDING_MODEL", "all-mpnet-base-v2") +st_model = os.environ.get("OPENEDAI_EMBEDDING_MODEL", "sentence-transformers/all-mpnet-base-v2") model = sentence_transformers.SentenceTransformer(st_model) diff --git a/extensions/openai/embeddings.py b/extensions/openai/embeddings.py index 396a3476..5482a3a5 100644 --- a/extensions/openai/embeddings.py +++ b/extensions/openai/embeddings.py @@ -3,8 +3,7 @@ import os import numpy as np from extensions.openai.errors import ServiceUnavailableError from extensions.openai.utils import debug_msg, float_list_to_base64 -from modules import shared -from transformers import AutoModel +from modules.logging_colors import logger embeddings_params_initialized = False @@ -16,38 +15,44 @@ def initialize_embedding_params(): ''' global embeddings_params_initialized if not embeddings_params_initialized: - global st_model, embeddings_model, embeddings_device from extensions.openai.script import params + + global st_model, embeddings_model, embeddings_device + st_model = os.environ.get("OPENEDAI_EMBEDDING_MODEL", params.get('embedding_model', 'all-mpnet-base-v2')) embeddings_model = None # OPENEDAI_EMBEDDING_DEVICE: auto (best or cpu), cpu, cuda, ipu, xpu, mkldnn, opengl, opencl, ideep, hip, ve, fpga, ort, xla, lazy, vulkan, mps, meta, hpu, mtia, privateuseone embeddings_device = os.environ.get("OPENEDAI_EMBEDDING_DEVICE", params.get('embedding_device', 'cpu')) if embeddings_device.lower() == 'auto': embeddings_device = None + embeddings_params_initialized = True def load_embedding_model(model: str): + try: + from sentence_transformers import SentenceTransformer + except ModuleNotFoundError: + logger.error("The sentence_transformers module has not been found. Please install it manually with pip install -U sentence-transformers.") + raise ModuleNotFoundError + initialize_embedding_params() global embeddings_device, embeddings_model try: print(f"Try embedding model: {model} on {embeddings_device}") - trust = shared.args.trust_remote_code - if embeddings_device == 'cpu': - embeddings_model = AutoModel.from_pretrained(model, trust_remote_code=trust).to("cpu", dtype=float) - else: #use the auto mode - embeddings_model = AutoModel.from_pretrained(model, trust_remote_code=trust) - print(f"\nLoaded embedding model: {model} on {embeddings_model.device}") + embeddings_model = SentenceTransformer(model, device=embeddings_device) + print(f"Loaded embedding model: {model}") except Exception as e: embeddings_model = None raise ServiceUnavailableError(f"Error: Failed to load embedding model: {model}", internal_message=repr(e)) -def get_embeddings_model() -> AutoModel: +def get_embeddings_model(): initialize_embedding_params() global embeddings_model, st_model if st_model and not embeddings_model: load_embedding_model(st_model) # lazy load the model + return embeddings_model @@ -66,9 +71,7 @@ def get_embeddings(input: list) -> np.ndarray: def embeddings(input: list, encoding_format: str) -> dict: - embeddings = get_embeddings(input) - if encoding_format == "base64": data = [{"object": "embedding", "embedding": float_list_to_base64(emb), "index": n} for n, emb in enumerate(embeddings)] else: @@ -85,5 +88,4 @@ def embeddings(input: list, encoding_format: str) -> dict: } debug_msg(f"Embeddings return size: {len(embeddings[0])}, number: {len(embeddings)}") - return response diff --git a/extensions/openai/models.py b/extensions/openai/models.py index 178ede90..c0a4606c 100644 --- a/extensions/openai/models.py +++ b/extensions/openai/models.py @@ -56,7 +56,7 @@ def _load_model(data): shared.model, shared.tokenizer = load_model(model_name) shared.model_name = model_name - + # Update shared.settings with custom generation defaults if settings: for k in settings: diff --git a/extensions/openai/script.py b/extensions/openai/script.py index 8bf78ca2..fd124b5f 100644 --- a/extensions/openai/script.py +++ b/extensions/openai/script.py @@ -31,6 +31,8 @@ from .typing import ( CompletionResponse, DecodeRequest, DecodeResponse, + EmbeddingsRequest, + EmbeddingsResponse, EncodeRequest, EncodeResponse, LoadModelRequest, @@ -41,7 +43,7 @@ from .typing import ( params = { 'embedding_device': 'cpu', - 'embedding_model': 'all-mpnet-base-v2', + 'embedding_model': 'sentence-transformers/all-mpnet-base-v2', 'sd_webui_url': '', 'debug': 0 } @@ -196,19 +198,16 @@ async def handle_image_generation(request: Request): return JSONResponse(response) -@app.post("/v1/embeddings") -async def handle_embeddings(request: Request): - body = await request.json() - encoding_format = body.get("encoding_format", "") - - input = body.get('input', body.get('text', '')) +@app.post("/v1/embeddings", response_model=EmbeddingsResponse) +async def handle_embeddings(request: Request, request_data: EmbeddingsRequest): + input = request_data.input if not input: raise HTTPException(status_code=400, detail="Missing required argument input") if type(input) is str: input = [input] - response = OAIembeddings.embeddings(input, encoding_format) + response = OAIembeddings.embeddings(input, request_data.encoding_format) return JSONResponse(response) diff --git a/extensions/openai/typing.py b/extensions/openai/typing.py index afd57eab..8b2938df 100644 --- a/extensions/openai/typing.py +++ b/extensions/openai/typing.py @@ -154,6 +154,19 @@ class LoadModelRequest(BaseModel): settings: dict | None = None +class EmbeddingsRequest(BaseModel): + input: str | List[str] + model: str | None = Field(default=None, description="Unused parameter. To change the model, set the OPENEDAI_EMBEDDING_MODEL and OPENEDAI_EMBEDDING_DEVICE environment variables before starting the server.") + encoding_format: str = Field(default="float", description="Can be float or base64.") + user: str | None = Field(default=None, description="Unused parameter.") + + +class EmbeddingsResponse(BaseModel): + index: int + embedding: List[float] + object: str = "embedding" + + def to_json(obj): return json.dumps(obj.__dict__, indent=4)