e

Dockerfile

@@ -17,4 +17,4 @@ WORKDIR $HOME/app
 
 COPY --chown=user . $HOME/app
 
-CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860", "--workers=4"]

main.py

@@ -1,11 +1,11 @@
 from typing import Generic, List, Optional, TypeVar
 from functools import partial
-from pydantic import BaseModel, ValidationError, validator
-from pydantic.generics import GenericModel
+from pydantic import BaseModel
 from sentence_transformers import SentenceTransformer
 from fastapi import FastAPI
-import os, asyncio, numpy, ujson
+import numpy
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import ORJSONResponse
 
 MODEL = SentenceTransformer("all-mpnet-base-v2")
 
@@ -24,24 +24,18 @@ def cache(func):
 
 @cache
 def _encode(sentences: List[str]):
-    array = [numpy.around(a.numpy(), 3) for a in MODEL.encode(sentences, normalize_embeddings=True, convert_to_tensor=True, batch_size=4, show_progress_bar=True)]
+    embeddings = MODEL.encode(sentences, normalize_embeddings=True, batch_size=2, show_progress_bar=True)
+    array = [numpy.around(a, 3).tolist() for a in embeddings]
     return array
 
-async def encode(sentences: List[str]) -> List[numpy.ndarray]:
-    loop = asyncio.get_event_loop()
-    result = await loop.run_in_executor(None, _encode, sentences)
-    return result
-
 class EmbedReq(BaseModel):
     sentences: List[str]
 
 app = FastAPI()
 
-@app.post("/embed")
-async def embed(embed: EmbedReq):
-    result = await encode(embed.sentences)
-    # Convert it to an ordinary list of floats
-    return ujson.dumps([r.tolist() for r in result])
+@app.post("/embed", response_class=ORJSONResponse)
+def embed(embed: EmbedReq):
+    return _encode(embed.sentences)
 
 app.add_middleware(
     CORSMiddleware,

requirements.txt

@@ -4,4 +4,4 @@ joblib==1.2
 fastapi==0.89
 uvicorn[standard]==0.20
 huggingface-hub==0.10.1
-ujson==5.7
+orjson==3.8.9