Setup for python environment.

Author: bertrama

py/Dockerfile

@@ -14,8 +14,8 @@ RUN mkdir -p /app
 WORKDIR /app
 
 ARG UNAME=app
-ARG UID
-ARG GID
+ARG UID=1000
+ARG GID=1000
 ENV APP_HOME /app
 
 #Create the group for the user
@@ -30,6 +30,6 @@ RUN if [ x"${UID}" != x"" ] ; \
     else useradd -m -d ${APP_HOME} -g ${UNAME} -s /bin/bash ${UNAME} ; \
     fi
 
-COPY requirements.txt .
+COPY requirements.txt sentence_embveddings.py setup.py startup.sh setup.sh /app/
 
-CMD tail -f /dev/null
+CMD /app/startup.sh

py/requirements.txt

@@ -8,11 +8,14 @@ sentence_transformers==2.2.2
 Pillow
 fastapi
 PyPDF2
-torch==2.0.1
+torch
 pydantic
 uvicorn
 sse-starlette
 boto3
 # missing from the langchain base image?
 langchain-openai
 langchain-community
+
+# https://github.com/easydiffusion/easydiffusion/issues/1851
+huggingface_hub==0.25.2

py/sentence_embeddings.py

@@ -0,0 +1,28 @@
+from fastapi import FastAPI, Depends, HTTPException, status, Query
+from pydantic import BaseModel
+from sentence_transformers import SentenceTransformer
+from typing import List, Optional
+import os
+
+app = FastAPI()
+models = {
+  'all-mpnet-base-v2': SentenceTransformer('all-mpnet-base-v2'),
+  'all-MiniLM-L6-v2': SentenceTransformer('all-MiniLM-L6-v2'),
+  'all-distilrobert-av1': SentenceTransformer('all-distilroberta-v1'),
+}
+
+class TextData(BaseModel):
+    texts: List[str]
+    model: str
+
+@app.post("/embeddings")
+async def create_embeddings(text_data: TextData):
+    try:
+        embeddings = models[text_data.model].encode(text_data.texts)
+        return {"embeddings": embeddings.tolist()}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=3000)

py/setup.py

@@ -0,0 +1,27 @@
+from fastapi import FastAPI, Depends, HTTPException, status, Query
+from pydantic import BaseModel
+from sentence_transformers import SentenceTransformer
+from typing import List, Optional
+import os
+
+app = FastAPI()
+models = {
+  'all-mpnet-base-v2': SentenceTransformer('all-mpnet-base-v2'),
+  'all-MiniLM-L6-v2': SentenceTransformer('all-MiniLM-L6-v2'),
+  'all-distilrobert-av1': SentenceTransformer('all-distilroberta-v1'),
+}
+
+class TextData(BaseModel):
+    texts: List[str]
+    model: str
+
+@app.post("/embeddings")
+async def create_embeddings(text_data: TextData):
+    try:
+        embeddings = models[text_data.model].encode(text_data.texts)
+        return {"embeddings": embeddings.tolist()}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    import uvicorn

py/setup.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+if [ ! -e /app/venv/bin/activate ] ; then
+  mkdir -p /app/venv/tmp /app/venv/cache/pip /app/lib/cache/hf
+  python -m venv /app/venv
+  source /app/venv/bin/activate
+  pip install -r requirements.txt
+  python setup.py
+fi

py/startup.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+source /app/venv/bin/activate
+python sentence_embeddings.py