-
-
Notifications
You must be signed in to change notification settings - Fork 2k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(transformers): split in separate backend
Signed-off-by: Ettore Di Giacinto <[email protected]>
- Loading branch information
Showing
23 changed files
with
777 additions
and
28 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,18 @@ | ||
| .PONY: sentencetransformers | ||
| sentencetransformers: | ||
| @echo "Creating virtual environment..." | ||
| @conda env create --name sentencetransformers --file sentencetransformers.yml | ||
| @echo "Virtual environment created." | ||
|
|
||
| .PONY: run | ||
| run: | ||
| @echo "Running sentencetransformers..." | ||
| bash run.sh | ||
| @echo "sentencetransformers run." | ||
|
|
||
| # It is not working well by using command line. It only6 works with IDE like VSCode. | ||
| .PONY: test | ||
| test: | ||
| @echo "Testing sentencetransformers..." | ||
| bash test.sh | ||
| @echo "sentencetransformers tested." |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,5 @@ | ||
| # Creating a separate environment for the sentencetransformers project | ||
|
|
||
| ``` | ||
| make sentencetransformers | ||
| ``` |
File renamed without changes.
File renamed without changes.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,14 @@ | ||
| #!/bin/bash | ||
|
|
||
| ## | ||
| ## A bash script wrapper that runs the sentencetransformers server with conda | ||
|
|
||
| export PATH=$PATH:/opt/conda/bin | ||
|
|
||
| # Activate conda environment | ||
| source activate sentencetransformers | ||
|
|
||
| # get the directory where the bash script is located | ||
| DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" | ||
|
|
||
| python $DIR/sentencetransformers.py $@ |
114 changes: 114 additions & 0 deletions
114
backend/python/sentencetransformers/sentencetransformers.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,114 @@ | ||
| """ | ||
| Extra gRPC server for HuggingFace SentenceTransformer models. | ||
| """ | ||
| #!/usr/bin/env python3 | ||
| from concurrent import futures | ||
|
|
||
| import argparse | ||
| import signal | ||
| import sys | ||
| import os | ||
|
|
||
| import time | ||
| import backend_pb2 | ||
| import backend_pb2_grpc | ||
|
|
||
| import grpc | ||
|
|
||
| from sentence_transformers import SentenceTransformer | ||
|
|
||
| _ONE_DAY_IN_SECONDS = 60 * 60 * 24 | ||
|
|
||
| # If MAX_WORKERS are specified in the environment use it, otherwise default to 1 | ||
| MAX_WORKERS = int(os.environ.get('PYTHON_GRPC_MAX_WORKERS', '1')) | ||
|
|
||
| # Implement the BackendServicer class with the service methods | ||
| class BackendServicer(backend_pb2_grpc.BackendServicer): | ||
| """ | ||
| A gRPC servicer for the backend service. | ||
| This class implements the gRPC methods for the backend service, including Health, LoadModel, and Embedding. | ||
| """ | ||
| def Health(self, request, context): | ||
| """ | ||
| A gRPC method that returns the health status of the backend service. | ||
| Args: | ||
| request: A HealthRequest object that contains the request parameters. | ||
| context: A grpc.ServicerContext object that provides information about the RPC. | ||
| Returns: | ||
| A Reply object that contains the health status of the backend service. | ||
| """ | ||
| return backend_pb2.Reply(message=bytes("OK", 'utf-8')) | ||
|
|
||
| def LoadModel(self, request, context): | ||
| """ | ||
| A gRPC method that loads a model into memory. | ||
| Args: | ||
| request: A LoadModelRequest object that contains the request parameters. | ||
| context: A grpc.ServicerContext object that provides information about the RPC. | ||
| Returns: | ||
| A Result object that contains the result of the LoadModel operation. | ||
| """ | ||
| model_name = request.Model | ||
| try: | ||
| self.model = SentenceTransformer(model_name) | ||
| except Exception as err: | ||
| return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}") | ||
|
|
||
| # Implement your logic here for the LoadModel service | ||
| # Replace this with your desired response | ||
| return backend_pb2.Result(message="Model loaded successfully", success=True) | ||
|
|
||
| def Embedding(self, request, context): | ||
| """ | ||
| A gRPC method that calculates embeddings for a given sentence. | ||
| Args: | ||
| request: An EmbeddingRequest object that contains the request parameters. | ||
| context: A grpc.ServicerContext object that provides information about the RPC. | ||
| Returns: | ||
| An EmbeddingResult object that contains the calculated embeddings. | ||
| """ | ||
| # Implement your logic here for the Embedding service | ||
| # Replace this with your desired response | ||
| print("Calculated embeddings for: " + request.Embeddings, file=sys.stderr) | ||
| sentence_embeddings = self.model.encode(request.Embeddings) | ||
| return backend_pb2.EmbeddingResult(embeddings=sentence_embeddings) | ||
|
|
||
|
|
||
| def serve(address): | ||
| server = grpc.server(futures.ThreadPoolExecutor(max_workers=MAX_WORKERS)) | ||
| backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server) | ||
| server.add_insecure_port(address) | ||
| server.start() | ||
| print("Server started. Listening on: " + address, file=sys.stderr) | ||
|
|
||
| # Define the signal handler function | ||
| def signal_handler(sig, frame): | ||
| print("Received termination signal. Shutting down...") | ||
| server.stop(0) | ||
| sys.exit(0) | ||
|
|
||
| # Set the signal handlers for SIGINT and SIGTERM | ||
| signal.signal(signal.SIGINT, signal_handler) | ||
| signal.signal(signal.SIGTERM, signal_handler) | ||
|
|
||
| try: | ||
| while True: | ||
| time.sleep(_ONE_DAY_IN_SECONDS) | ||
| except KeyboardInterrupt: | ||
| server.stop(0) | ||
|
|
||
| if __name__ == "__main__": | ||
| parser = argparse.ArgumentParser(description="Run the gRPC server.") | ||
| parser.add_argument( | ||
| "--addr", default="localhost:50051", help="The address to bind the server to." | ||
| ) | ||
| args = parser.parse_args() | ||
|
|
||
| serve(args.addr) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,11 @@ | ||
| #!/bin/bash | ||
| ## | ||
| ## A bash script wrapper that runs the sentencetransformers server with conda | ||
|
|
||
| # Activate conda environment | ||
| source activate sentencetransformers | ||
|
|
||
| # get the directory where the bash script is located | ||
| DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" | ||
|
|
||
| python -m unittest $DIR/test_sentencetransformers.py |
81 changes: 81 additions & 0 deletions
81
backend/python/sentencetransformers/test_sentencetransformers.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,81 @@ | ||
| """ | ||
| A test script to test the gRPC service | ||
| """ | ||
| import unittest | ||
| import subprocess | ||
| import time | ||
| import backend_pb2 | ||
| import backend_pb2_grpc | ||
|
|
||
| import grpc | ||
|
|
||
|
|
||
| class TestBackendServicer(unittest.TestCase): | ||
| """ | ||
| TestBackendServicer is the class that tests the gRPC service | ||
| """ | ||
| def setUp(self): | ||
| """ | ||
| This method sets up the gRPC service by starting the server | ||
| """ | ||
| self.service = subprocess.Popen(["python3", "sentencetransformers.py", "--addr", "localhost:50051"]) | ||
|
|
||
| def tearDown(self) -> None: | ||
| """ | ||
| This method tears down the gRPC service by terminating the server | ||
| """ | ||
| self.service.terminate() | ||
| self.service.wait() | ||
|
|
||
| def test_server_startup(self): | ||
| """ | ||
| This method tests if the server starts up successfully | ||
| """ | ||
| time.sleep(2) | ||
| try: | ||
| self.setUp() | ||
| with grpc.insecure_channel("localhost:50051") as channel: | ||
| stub = backend_pb2_grpc.BackendStub(channel) | ||
| response = stub.Health(backend_pb2.HealthMessage()) | ||
| self.assertEqual(response.message, b'OK') | ||
| except Exception as err: | ||
| print(err) | ||
| self.fail("Server failed to start") | ||
| finally: | ||
| self.tearDown() | ||
|
|
||
| def test_load_model(self): | ||
| """ | ||
| This method tests if the model is loaded successfully | ||
| """ | ||
| try: | ||
| self.setUp() | ||
| with grpc.insecure_channel("localhost:50051") as channel: | ||
| stub = backend_pb2_grpc.BackendStub(channel) | ||
| response = stub.LoadModel(backend_pb2.ModelOptions(Model="bert-base-nli-mean-tokens")) | ||
| self.assertTrue(response.success) | ||
| self.assertEqual(response.message, "Model loaded successfully") | ||
| except Exception as err: | ||
| print(err) | ||
| self.fail("LoadModel service failed") | ||
| finally: | ||
| self.tearDown() | ||
|
|
||
| def test_embedding(self): | ||
| """ | ||
| This method tests if the embeddings are generated successfully | ||
| """ | ||
| try: | ||
| self.setUp() | ||
| with grpc.insecure_channel("localhost:50051") as channel: | ||
| stub = backend_pb2_grpc.BackendStub(channel) | ||
| response = stub.LoadModel(backend_pb2.ModelOptions(Model="bert-base-nli-mean-tokens")) | ||
| self.assertTrue(response.success) | ||
| embedding_request = backend_pb2.PredictOptions(Embeddings="This is a test sentence.") | ||
| embedding_response = stub.Embedding(embedding_request) | ||
| self.assertIsNotNone(embedding_response.embeddings) | ||
| except Exception as err: | ||
| print(err) | ||
| self.fail("Embedding service failed") | ||
| finally: | ||
| self.tearDown() |
14 changes: 7 additions & 7 deletions
14
backend/python/huggingface/Makefile → backend/python/transformers/Makefile
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,18 +1,18 @@ | ||
| .PONY: huggingface | ||
| huggingface: | ||
| .PONY: transformers | ||
| transformers: | ||
| @echo "Creating virtual environment..." | ||
| @conda env create --name huggingface --file huggingface.yml | ||
| @conda env create --name transformers --file transformers.yml | ||
| @echo "Virtual environment created." | ||
|
|
||
| .PONY: run | ||
| run: | ||
| @echo "Running huggingface..." | ||
| @echo "Running transformers..." | ||
| bash run.sh | ||
| @echo "huggingface run." | ||
| @echo "transformers run." | ||
|
|
||
| # It is not working well by using command line. It only6 works with IDE like VSCode. | ||
| .PONY: test | ||
| test: | ||
| @echo "Testing huggingface..." | ||
| @echo "Testing transformers..." | ||
| bash test.sh | ||
| @echo "huggingface tested." | ||
| @echo "transformers tested." |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,5 @@ | ||
| # Creating a separate environment for the transformers project | ||
|
|
||
| ``` | ||
| make transformers | ||
| ``` |
Oops, something went wrong.