Support for text-embeddings-inference

Author: xusenlin

api/config.py

@@ -39,7 +39,7 @@ class Settings(BaseModel):
  )
  engine: Optional[str] = Field(
  default=get_env("ENGINE", "default"),
- description="Choices are ['default', 'vllm', 'llama.cpp'].",
+ description="Choices are ['default', 'vllm', 'llama.cpp', 'tgi'].",
  )
 
  # model related
@@ -239,10 +239,24 @@ class Settings(BaseModel):
  description="RoPE frequency scaling factor",
  )
 
- # support for tgi
+ # support for tgi: https://github.com/huggingface/text-generation-inference
  tgi_endpoint: Optional[str] = Field(
  default=get_env("TGI_ENDPOINT", None),
- description="Text Generate Inference Endpoint.",
+ description="Text Generation Inference Endpoint.",
+ )
+
+ # support for tei: https://github.com/huggingface/text-embeddings-inference
+ tei_endpoint: Optional[str] = Field(
+ default=get_env("TEI_ENDPOINT", None),
+ description="Text Embeddings Inference Endpoint.",
+ )
+ max_concurrent_requests: Optional[int] = Field(
+ default=int(get_env("MAX_CONCURRENT_REQUESTS", 256)),
+ description="The maximum amount of concurrent requests for this particular deployment."
+ )
+ max_client_batch_size: Optional[int] = Field(
+ default=int(get_env("MAX_CLIENT_BATCH_SIZE", 32)),
+ description="Control the maximum number of inputs that a client can send in a single request."
  )
 
 

api/models.py

@@ -21,9 +21,13 @@ def create_app() -> FastAPI:
 
 def create_embedding_model():
  """ get embedding model from sentence-transformers. """
- from sentence_transformers import SentenceTransformer
-
- return SentenceTransformer(SETTINGS.embedding_name, device=SETTINGS.embedding_device)
+ if SETTINGS.tei_endpoint is not None:
+ from openai import AsyncOpenAI
+ client = AsyncOpenAI(base_url=SETTINGS.tei_endpoint, api_key="none")
+ else:
+ from sentence_transformers import SentenceTransformer
+ client = SentenceTransformer(SETTINGS.embedding_name, device=SETTINGS.embedding_device)
+ return client
 
 
 def create_generate_model():

api/routes/embedding.py

@@ -1,8 +1,11 @@
+import asyncio
 import base64
+from typing import Union
 
 import numpy as np
 import tiktoken
 from fastapi import APIRouter, Depends
+from openai import AsyncOpenAI
 from openai.types.create_embedding_response import Usage
 from sentence_transformers import SentenceTransformer
 
@@ -23,7 +26,7 @@ def get_embedding_engine():
 async def create_embeddings(
  request: EmbeddingCreateParams,
  model_name: str = None,
- engine: SentenceTransformer = Depends(get_embedding_engine),
+ client: Union[SentenceTransformer, AsyncOpenAI] = Depends(get_embedding_engine),
 ):
  """Creates embeddings for the text"""
  if request.model is None:
@@ -41,7 +44,7 @@ async def create_embeddings(
  request.input = [decoding.decode(text) for text in request.input]
 
  # https://huggingface.co/BAAI/bge-large-zh
- if engine is not None and "bge" in SETTINGS.embedding_name.lower():
+ if client is not None and "bge" in SETTINGS.embedding_name.lower():
  instruction = ""
  if "zh" in SETTINGS.embedding_name.lower():
  instruction = "为这个句子生成表示以用于检索相关文章："
@@ -50,30 +53,65 @@ async def create_embeddings(
  request.input = [instruction + q for q in request.input]
 
  data, total_tokens = [], 0
- batches = [
- request.input[i: i + 1024] for i in range(0, len(request.input), 1024)
- ]
- for num_batch, batch in enumerate(batches):
- token_num = sum(len(i) for i in batch)
- vecs = engine.encode(batch, normalize_embeddings=True)
-
- bs, dim = vecs.shape
- if SETTINGS.embedding_size > dim:
- zeros = np.zeros((bs, SETTINGS.embedding_size - dim))
- vecs = np.c_[vecs, zeros]
-
- if request.encoding_format == "base64":
- vecs = [base64.b64encode(v.tobytes()).decode("utf-8") for v in vecs]
- else:
- vecs = vecs.tolist()
-
- data.extend(
- Embedding(
- index=num_batch * 1024 + i, object="embedding", embedding=embed
+
+ # support for tei: https://github.com/huggingface/text-embeddings-inference
+ if isinstance(client, AsyncOpenAI):
+ global_batch_size = SETTINGS.max_concurrent_requests * SETTINGS.max_client_batch_size
+ for i in range(0, len(request.input), global_batch_size):
+ tasks = []
+ texts = request.input[i: i + global_batch_size]
+ for j in range(0, len(texts), SETTINGS.max_client_batch_size):
+ tasks.append(
+ client.embeddings.create(
+ input=texts[j: j + SETTINGS.max_client_batch_size],
+ model=request.model,
+ )
+ )
+ res = await asyncio.gather(*tasks)
+
+ vecs = np.asarray([e.embedding for r in res for e in r.data])
+ bs, dim = vecs.shape
+ if SETTINGS.embedding_size > dim:
+ zeros = np.zeros((bs, SETTINGS.embedding_size - dim))
+ vecs = np.c_[vecs, zeros]
+
+ if request.encoding_format == "base64":
+ vecs = [base64.b64encode(v.tobytes()).decode("utf-8") for v in vecs]
+ else:
+ vecs = vecs.tolist()
+
+ data.extend(
+ Embedding(
+ index=i * global_batch_size + j,
+ object="embedding",
+ embedding=embed
+ )
+ for j, embed in enumerate(vecs)
+ )
+ total_tokens += sum(r.usage.total_tokens for r in res)
+ else:
+ batches = [request.input[i: i + 1024] for i in range(0, len(request.input), 1024)]
+ for num_batch, batch in enumerate(batches):
+ token_num = sum(len(i) for i in batch)
+ vecs = client.encode(batch, normalize_embeddings=True)
+
+ bs, dim = vecs.shape
+ if SETTINGS.embedding_size > dim:
+ zeros = np.zeros((bs, SETTINGS.embedding_size - dim))
+ vecs = np.c_[vecs, zeros]
+
+ if request.encoding_format == "base64":
+ vecs = [base64.b64encode(v.tobytes()).decode("utf-8") for v in vecs]
+ else:
+ vecs = vecs.tolist()
+
+ data.extend(
+ Embedding(
+ index=num_batch * 1024 + i, object="embedding", embedding=embed
+ )
+ for i, embed in enumerate(vecs)
  )
- for i, embed in enumerate(vecs)
- )
- total_tokens += token_num
+ total_tokens += token_num
 
  return CreateEmbeddingResponse(
  data=data,

api/routes/model.py

@@ -19,7 +19,7 @@ class ModelList(BaseModel):
 available_models = ModelList(
  data=[
  Model(
- id=SETTINGS.model_name,
+ id=SETTINGS.model_name or "",
  object="model",
  created=int(time.time()),
  owned_by="open"

api/server.py

@@ -1,10 +1,8 @@
 from api.config import SETTINGS
 from api.models import app, EMBEDDED_MODEL, GENERATE_ENGINE
-from api.routes import model_router
 
 
 prefix = SETTINGS.api_prefix
-app.include_router(model_router, prefix=prefix, tags=["Model"])
 
 if EMBEDDED_MODEL is not None:
  from api.routes.embedding import embedding_router
@@ -13,6 +11,10 @@
 
 
 if GENERATE_ENGINE is not None:
+ from api.routes import model_router
+
+ app.include_router(model_router, prefix=prefix, tags=["Model"])
+
  if SETTINGS.engine == "vllm":
  from api.vllm_routes import chat_router as chat_router
  from api.vllm_routes import completion_router as completion_router
@@ -29,7 +31,7 @@
  from api.routes.chat import chat_router as chat_router
  from api.routes.completion import completion_router as completion_router
 
- app.include_router(chat_router, prefix=prefix, tags=["Chat"])
+ app.include_router(chat_router, prefix=prefix, tags=["Chat Completion"])
  app.include_router(completion_router, prefix=prefix, tags=["Completion"])