Add support for reranker-transformers module

weaviate · Jul 6, 2023 · adcc5a2 · adcc5a2
1 parent 8b2fc03
commit adcc5a2
Show file tree

Hide file tree

Showing 6 changed files with 164 additions and 2 deletions.
diff --git a/.cicd/test.sh b/.cicd/test.sh
@@ -140,6 +140,8 @@ function check_creates_template() {
   check_modules "--set modules.generative-palm.enabled=true" "value: generative-palm"
   check_modules "--set modules.text2vec-palm.enabled=true --set modules.text2vec-palm.apiKey=apiKey --set modules.generative-palm.enabled=true --set modules.generative-palm.apiKey=apiKey" "value: generative-palm,text2vec-palm"
   check_modules "--set modules.generative-palm.enabled=true --set modules.generative-palm.apiKey=apiKey" "value: generative-palm"
+  check_modules "--set modules.text2vec-contextionary.enabled=false --set modules.reranker-transformers.enabled=true" "value: reranker-transformers"
+  check_modules "--set modules.text2vec-contextionary.enabled=true --set modules.text-spellcheck.enabled=true --set modules.reranker-transformers.enabled=true" "value: text2vec-contextionary,text-spellcheck,reranker-transformers"
 
   _settingPassageQueryOn="--set modules.text2vec-contextionary.enabled=false --set modules.text2vec-transformers.passageQueryServices.passage.enabled=true --set modules.text2vec-transformers.passageQueryServices.query.enabled=true"
   check_setting_has_value "$_settingPassageQueryOn" "name: TRANSFORMERS_PASSAGE_INFERENCE_API" "value: http://transformers-inference-passage.default.svc.cluster.local:8080"
@@ -164,7 +166,7 @@ function check_creates_template() {
   check_string_existence "--set initContainers.sysctlInitContainer.enabled=false --set initContainers.extraInitContainers[0].name=test-init-container " "name: test-init-container"
 
   check_string_existence "" "imagePullPolicy: IfNotPresent"
-  MODULES=("text2vec-transformers" "multi2vec-clip" "qna-transformers" "img2vec-neural" "text-spellcheck" "ner-transformers" "sum-transformers")
+  MODULES=("text2vec-transformers" "multi2vec-clip" "qna-transformers" "img2vec-neural" "text-spellcheck" "ner-transformers" "sum-transformers" "reranker-transformers")
   for mod in "${MODULES[@]}"
   do
   check_string_existence "--set modules.$mod.enabled=true --set modules.$mod.livenessProbe.initialDelaySeconds=988888888888" "initialDelaySeconds: 988888888888"
@@ -187,7 +189,7 @@ function check_creates_template() {
 
   check_setting_has_value "--set image.pullSecrets[0]=weaviate-image-pull-secret" "imagePullSecrets" "name: weaviate-image-pull-secret"
 
-  MODULES=("text2vec-contextionary" "text2vec-transformers" "multi2vec-clip" "qna-transformers" "img2vec-neural" "text-spellcheck" "ner-transformers" "sum-transformers")
+  MODULES=("text2vec-contextionary" "text2vec-transformers" "multi2vec-clip" "qna-transformers" "img2vec-neural" "text-spellcheck" "ner-transformers" "sum-transformers" "reranker-transformers")
   for mod in "${MODULES[@]}"
   do
   check_setting_has_value "--set modules.$mod.enabled=true --set modules.$mod.imagePullSecrets[0]=weaviate-image-pull-secret" "imagePullSecrets" "name: weaviate-image-pull-secret"

diff --git a/weaviate/templates/_helpers.tpl b/weaviate/templates/_helpers.tpl
@@ -59,6 +59,9 @@
   {{- if (index .Values "modules" "reranker-cohere" "enabled") -}}
     {{ $modules = append $modules "reranker-cohere" }}
   {{- end -}}
+  {{- if (index .Values "modules" "reranker-transformers" "enabled") -}}
+    {{ $modules = append $modules "reranker-transformers" }}
+  {{- end -}}
   {{- if (index .Values "backups" "filesystem" "enabled") -}}
     {{ $modules = append $modules "backup-filesystem" }}
   {{- end -}}

diff --git a/weaviate/templates/rerankerTransformersDeployment.yaml b/weaviate/templates/rerankerTransformersDeployment.yaml
@@ -0,0 +1,69 @@
+{{ if index .Values "modules" "reranker-transformers" "enabled" }}
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: {{ index .Values "modules" "reranker-transformers" "fullnameOverride" }}
+  labels:
+    name: {{ index .Values "modules" "reranker-transformers" "fullnameOverride" }}
+    app: {{ index .Values "modules" "reranker-transformers" "fullnameOverride" }}
+    app.kubernetes.io/name: weaviate
+    app.kubernetes.io/managed-by: helm
+spec:
+  replicas: {{ index .Values "modules" "reranker-transformers" "replicas" }}
+  selector:
+    matchLabels:
+      app: {{ index .Values "modules" "reranker-transformers" "fullnameOverride" }}
+  template:
+    metadata:
+      labels:
+        app: {{ index .Values "modules" "reranker-transformers" "fullnameOverride" }}
+    spec:
+      {{- if (index .Values "modules" "reranker-transformers" "serviceAccountName") }}
+      serviceAccountName: {{ index .Values "modules" "reranker-transformers" "serviceAccountName" }}
+      {{- else if (index .Values "serviceAccountName") }}
+      serviceAccountName: {{ index .Values "serviceAccountName" }}
+      {{- end }}
+      {{- include "image.pullSecrets" (dict "pullSecrets" (index .Values "modules" "reranker-transformers" "imagePullSecrets")) | nindent 6 }}
+      containers:
+      - name: {{ index .Values "modules" "reranker-transformers" "fullnameOverride" }}
+        image: {{ index .Values "modules" "reranker-transformers" "registry" }}/{{ index .Values "modules" "reranker-transformers" "repo" }}:{{ index .Values "modules" "reranker-transformers" "tag" }}
+        imagePullPolicy: {{ index .Values "modules" "reranker-transformers" "imagePullPolicy" }}
+        env:
+          - name: ENABLE_CUDA
+            value: "{{ if index .Values "modules" "reranker-transformers" "envconfig" "enable_cuda" }}1{{ else }}0{{ end }}"
+          {{ if index .Values "modules" "reranker-transformers" "envconfig" "enable_cuda" }}
+          - name: NVIDIA_VISIBLE_DEVICES
+            value: {{ index .Values "modules" "reranker-transformers" "envconfig" "nvidia_visible_devices"}}
+          - name: NVIDIA_DRIVER_CAPABILITIES
+            value: {{ index .Values "modules" "reranker-transformers" "envconfig" "nvidia_driver_capabilities"}}
+          - name: LD_LIBRARY_PATH
+            value: {{ index .Values "modules" "reranker-transformers" "envconfig" "ld_library_path"}}
+          {{ end }}
+        resources:
+{{ index .Values "modules" "reranker-transformers" "resources" | toYaml | indent 10 }}
+        livenessProbe:
+          httpGet:
+            path: /.well-known/live
+            port: 8080
+          initialDelaySeconds: {{ index .Values "modules" "reranker-transformers" "livenessProbe" "initialDelaySeconds" }}
+          periodSeconds: {{ index .Values "modules" "reranker-transformers" "livenessProbe" "periodSeconds" }}
+          timeoutSeconds: {{ index .Values "modules" "reranker-transformers" "livenessProbe" "timeoutSeconds" }}
+        readinessProbe:
+          httpGet:
+            path: /.well-known/ready
+            port: 8080
+          initialDelaySeconds: {{ index .Values "modules" "reranker-transformers" "readinessProbe" "initialDelaySeconds" }}
+          periodSeconds: {{ index .Values "modules" "reranker-transformers" "readinessProbe" "periodSeconds" }}
+      {{- with index .Values "modules" "reranker-transformers" "nodeSelector" | default .Values.nodeSelector }}
+      nodeSelector:
+        {{ toYaml . | nindent 8 }}
+      {{- end }}
+      {{- with index .Values "modules" "reranker-transformers" "affinity" | default .Values.affinity }}
+      affinity:
+        {{ toYaml . | nindent 8 }}
+      {{- end }}
+      {{- with index .Values "modules" "reranker-transformers" "tolerations" | default .Values.tolerations }}
+      tolerations:
+        {{ toYaml . | nindent 8 }}
+      {{- end }}
+{{ end }}
diff --git a/weaviate/templates/rerankerTransformersService.yaml b/weaviate/templates/rerankerTransformersService.yaml
@@ -0,0 +1,17 @@
+{{ if index .Values "modules" "reranker-transformers" "enabled" }}
+apiVersion: v1
+kind: Service
+metadata:
+  name: {{ index .Values "modules" "reranker-transformers" "fullnameOverride" }}
+  labels:
+    app.kubernetes.io/name: weaviate
+    app.kubernetes.io/managed-by: helm
+spec:
+  type: ClusterIP
+  selector:
+    app: {{ index .Values "modules" "reranker-transformers" "fullnameOverride" }}
+  ports:
+    - protocol: TCP
+      port: 8080
+      targetPort: 8080
+{{ end }}
diff --git a/weaviate/templates/weaviateStatefulset.yaml b/weaviate/templates/weaviateStatefulset.yaml
@@ -147,6 +147,13 @@ spec:
           - name: SUM_INFERENCE_API
             value: {{ index .Values "modules" "sum-transformers" "inferenceUrl" }}
           {{- end }}
+          {{- if index .Values "modules" "reranker-transformers" "enabled" }}
+          - name: RERANKER_INFERENCE_API
+            value: http://{{ index .Values "modules" "reranker-transformers" "fullnameOverride" }}.{{ .Release.Namespace }}.svc.cluster.local:8080
+          {{- else if index .Values "modules" "reranker-transformers" "inferenceUrl" }}
+          - name: RERANKER_INFERENCE_API
+            value: {{ index .Values "modules" "reranker-transformers" "inferenceUrl" }}
+          {{- end }}
           {{ if or (and (index .Values "modules" "text2vec-openai" "enabled") (index .Values "modules" "text2vec-openai" "apiKey")) (and (index .Values "modules" "qna-openai" "enabled") (index .Values "modules" "qna-openai" "apiKey")) (and (index .Values "modules" "generative-openai" "enabled") (index .Values "modules" "generative-openai" "apiKey")) }}
           - name: OPENAI_APIKEY
             valueFrom:

diff --git a/weaviate/values.yaml b/weaviate/values.yaml
@@ -893,6 +893,70 @@ modules:
     # an environment variable
     apiKey: ''
 
+  # The reranker-transformers module uses Cross-Encoders for 
+  # sentence pair scoring and sentence pair classification tasks.
+  # More information about Cross-Encoders can be found here:
+  # https://www.sbert.net/examples/applications/cross-encoder/README.html
+  reranker-transformers:
+    enabled: false
+    # You can set directly an inference URL of this module without deploying it with this release.
+    # You can do so by setting a value for the `inferenceUrl` here AND by setting the `enable` to `false`
+    inferenceUrl: {}
+    tag: cross-encoder-ms-marco-MiniLM-L-6-v2
+    repo: semitechnologies/reranker-transformers
+    registry: docker.io
+    replicas: 1
+    imagePullPolicy: IfNotPresent
+    imagePullSecrets: []
+    fullnameOverride: reranker-transformers
+    livenessProbe:
+      initialDelaySeconds: 120
+      periodSeconds: 3
+      timeoutSeconds: 3
+    readinessProbe:
+      initialDelaySeconds: 120
+      periodSeconds: 3
+    envconfig:
+      # enable for CUDA support. Your K8s cluster needs to be configured
+      # accordingly and you need to explicitly set GPU requests & limits below
+      enable_cuda: false
+
+      # only used when cuda is enabled
+      nvidia_visible_devices: all
+      nvidia_driver_capabilities: compute,utility
+
+      # only used when cuda is enabled
+      ld_library_path: /usr/local/nvidia/lib64
+
+    resources:
+      requests:
+        cpu: '1000m'
+        memory: '3000Mi'
+
+        # enable if running with CUDA support
+        # nvidia.com/gpu: 1
+      limits:
+        cpu: '1000m'
+        memory: '5000Mi'
+
+        # enable if running with CUDA support
+        # nvidia.com/gpu: 1
+
+    # It is possible to add a ServiceAccount to this module's Pods, it can be
+    # used in cases where the module is in a private registry and you want to
+    # give access to the registry only to this pod.
+    # NOTE: if not set the root `serviceAccountName` config will be used.
+    serviceAccountName:
+
+    # You can guide where the pods are scheduled on a per-module basis,
+    # as well as for Weaviate overall. Each module accepts nodeSelector,
+    # tolerations, and affinity configuration. If it is set on a per-
+    # module basis, this configuration overrides the global config.
+
+    nodeSelector:
+    tolerations:
+    affinity:
+
   # The text-spellcheck module uses spellchecker library to check
   # misspellings in a given text
   text-spellcheck: