Adding example for sentence embedding

Author: lno

examples/sentence-transformers/Dockerfile

@@ -0,0 +1,29 @@
+FROM nvidia/cuda:11.7.1-devel-ubuntu20.04
+
+# Update, install
+RUN apt update && \
+    apt install -y build-essential python3 python3-pip git
+
+RUN pip install --upgrade pip setuptools wheel
+
+# Create user instead of using root
+ENV USER='user'
+RUN groupadd -r user && useradd -r -g $USER $USER
+USER $USER
+
+# Define workdir
+WORKDIR /home/$USER/app
+
+# Install project
+COPY . .
+RUN pip install -r requirements.txt
+
+# Get model weights and tokenizer
+RUN python3 get_models.py
+
+# Publish port
+EXPOSE 50051:50051
+
+# Enjoy
+ENTRYPOINT ["python3", "server.py"]
+CMD ["--address", "[::]:50051"]

examples/sentence-transformers/README.md

@@ -0,0 +1,22 @@
+# OpenChatKit service
+
+## Description
+
+This project use the [OpenChatKit](https://github.com/togethercomputer/OpenChatKit) model and turns it into an gRPC service that can be used through [SimpleAI](https://github.com/lhenault/simpleAI).
+
+To quote the project:
+
+>
+    OpenChatKit provides a powerful, open-source base to create both specialized and general purpose chatbots for various applications. The kit includes an instruction-tuned 20 billion parameter language model, a 6 billion parameter moderation model, and an extensible retrieval system for including up-to-date responses from custom repositories. It was trained on the OIG-43M training dataset, which was a collaboration between Together, LAION, and Ontocord.ai. Much more than a model release, this is the beginning of an open source project. We are releasing a set of tools and processes for ongoing improvement with community contributions.
+
+## Build
+
+```bash
+docker build . -t openchatkit-service:latest
+```
+
+## Start service
+
+```bash
+docker run -it --rm --p 50051:50051 --gpus all openchatkit-service:latest
+```

examples/sentence-transformers/get_models.py

@@ -0,0 +1,11 @@
+import logging
+
+from sentence_transformers import SentenceTransformer
+
+MODEL_ID = 'all-MiniLM-L6-v2'
+
+if __name__ == '__main__':
+    try:
+        model = SentenceTransformer(MODEL_ID)
+    except Exception as ex:
+        logging.exception(f'Could not retrieve {MODEL_ID}: {ex}')

examples/sentence-transformers/model.py

@@ -0,0 +1,19 @@
+import logging
+from dataclasses import dataclass
+
+from sentence_transformers import SentenceTransformer
+from simple_ai.serve.python.embedding.model import LanguageModel
+
+
+from get_models import MODEL_ID
+
+@dataclass(unsafe_hash=True)
+class SentenceTransformerModel:
+    model = SentenceTransformer(MODEL_ID)
+    def embed(self, 
+        inputs: list=[],
+    ) -> list:
+        logging.info(f'Processing inputs : {inputs}')
+        embeddings = self.model.encode(inputs)
+        logging.info(f'Successfully computed embeddings (shape : {embeddings.shape}) for inputs : {inputs}')
+        return embeddings.tolist()

examples/sentence-transformers/requirements.txt

@@ -0,0 +1,2 @@
+sentence-transformers
+simple_ai_server>=0.1.9

examples/sentence-transformers/server.py

@@ -0,0 +1,20 @@
+import logging
+
+from simple_ai.serve.python.embedding.server import serve, LanguageModelServicer
+
+from model import SentenceTransformerModel as Model
+
+
+if __name__ == '__main__':
+    import argparse
+    
+    logging.basicConfig(level=logging.INFO)
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-a', '--address', type=str, default='[::]:50051')
+    args = parser.parse_args()
+
+    logging.info(f'Starting gRPC server on {args.address}')
+    
+    model_servicer = LanguageModelServicer(model=Model())
+    serve(address=args.address, model_servicer=model_servicer)