v1 with gradio + custom api endpoint

Dockerfile

@@ -7,7 +7,7 @@ ARG DEBIAN_FRONTEND=noninteractive
 
 RUN apt-get update
 
-RUN python -m pip install spaces
+RUN python -m pip install spaces pydantic
 
 WORKDIR /app
 COPY requirements.txt .

app.py

@@ -1,20 +1,58 @@
 import gradio as gr
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
 import uvicorn
+import spaces
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.util import cos_sim
+from sentence_transformers.quantization import quantize_embeddings
+
 
 app = FastAPI()
 
+print("Loading embedding model");
+Embedder = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
+
+@spaces.GPU
+def embed(text):
+        
+    query_embedding = Embedder.encode(text)
+    return query_embedding.tolist();
+    
+    
+
 
 with gr.Blocks(fill_height=True) as demo:
-       text = gr.Textbox();
+    text = gr.Textbox();
+    embeddings = gr.Textbox()
+    
+    text.submit(embed, [text], [embeddings]);
+    
 
 
 
 
 
-@app.get("/v1/embeddings")
-def test():
-    return {"test": True}
+@app.post("/v1/embeddings")
+async def openai_embeddings(request: Request):
+    body = await request.json();
+    print(body);
+    
+    model = body['model']
+    text = body['input'];
+    embeddings = embed(text)
+    return {
+		'object': "list"
+		,'data': [{
+			'object': "embeddings"
+			,'embedding': embeddings
+			,'index':0
+		}]
+		,'model':model
+		,'usage':{
+			 'prompt_tokens': 0
+			,'total_tokens': 0
+		}
+	}
     
 
 

requirements.txt

@@ -1,2 +1,3 @@
 fastapi
-uvicorn
+uvicorn
+sentence_transformers