From 4a4e44bf5559f2eac49df2c1135f39ad6d70300f Mon Sep 17 00:00:00 2001
From: Alexander Izotov <93216976+Nyralei@users.noreply.github.com>
Date: Tue, 17 Sep 2024 08:52:37 +0300
Subject: [PATCH] feat: allow setting trust_remote_code for
 sentencetransformers backend (#3552)

Allow setting trust_remote_code for SentenceTransformers backend

Signed-off-by: Nyralei <93216976+Nyralei@users.noreply.github.com>
---
 backend/python/sentencetransformers/backend.py       | 2 +-
 backend/python/sentencetransformers/requirements.txt | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/backend/python/sentencetransformers/backend.py b/backend/python/sentencetransformers/backend.py
index 905015e1..2a20bf60 100755
--- a/backend/python/sentencetransformers/backend.py
+++ b/backend/python/sentencetransformers/backend.py
@@ -55,7 +55,7 @@ class BackendServicer(backend_pb2_grpc.BackendServicer):
         """
         model_name = request.Model
         try:
-            self.model = SentenceTransformer(model_name)
+            self.model = SentenceTransformer(model_name, trust_remote_code=request.TrustRemoteCode)
         except Exception as err:
             return backend_pb2.Result(success=False, message=f"Unexpected {err=}, {type(err)=}")
 
diff --git a/backend/python/sentencetransformers/requirements.txt b/backend/python/sentencetransformers/requirements.txt
index 8e1b0195..b9cb6061 100644
--- a/backend/python/sentencetransformers/requirements.txt
+++ b/backend/python/sentencetransformers/requirements.txt
@@ -1,3 +1,5 @@
 grpcio==1.66.1
 protobuf
-certifi
\ No newline at end of file
+certifi
+datasets
+einops
\ No newline at end of file