Browse Source

feat(template): add TEI template

Thomas Zhang 2 months ago
parent
commit
ac4ff05906
2 changed files with 85 additions and 0 deletions
  1. 81 0
      pkg/templates/tei/template.go
  2. 4 0
      pkg/templates/templates.go

+ 81 - 0
pkg/templates/tei/template.go

@@ -0,0 +1,81 @@
+/*
+Copyright 2026 LocoStack.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package tei
+
+import (
+	"github.com/LocoStack/loco-operator/api/v1alpha1"
+	corev1 "k8s.io/api/core/v1"
+)
+
+const (
+	TEI_MODEL_PATH = "/models"
+)
+
+var TEITemplate = v1alpha1.Template{
+	Name: "tei",
+	Spec: v1alpha1.TemplateSpec{
+		Runtime: v1alpha1.RuntimeSpec{
+			Name:  "tei",
+			Image: "ghcr.io/huggingface/text-embeddings-inference:cpu-latest",
+			Args: []string{
+				"--model-id", TEI_MODEL_PATH,
+				"--port", "$(PORT)",
+				"--max-batch-tokens", "$(spec.runtimeInferenceParameters.contextWindow)",
+				"--max-client-batch-size", "8",
+				"--auto-truncate",
+			},
+			Env: []corev1.EnvVar{
+				{Name: "MODEL_NAME", Value: "$(spec.modelName)"},
+				{Name: "PORT", Value: "$(spec.runtime.port)"},
+			},
+			VolumeMounts: []corev1.VolumeMount{
+				{Name: "model", MountPath: TEI_MODEL_PATH},
+			},
+			Port: 80,
+		},
+		Volumes: []corev1.Volume{
+			{
+				Name: "model",
+				VolumeSource: corev1.VolumeSource{
+					PersistentVolumeClaim: &corev1.PersistentVolumeClaimVolumeSource{
+						ClaimName: "$(spec.pvcName)",
+					},
+				},
+			},
+		},
+	},
+}
+
+var TEIARM64Runtime = v1alpha1.RuntimeSpec{
+	Name:  "tei-arm64",
+	Image: "ghcr.io/huggingface/text-embeddings-inference:cpu-arm64-latest",
+	Args: []string{
+		"--model-id", TEI_MODEL_PATH,
+		"--port", "$(PORT)",
+		"--max-batch-tokens", "$(spec.runtimeInferenceParameters.contextWindow)",
+		"--max-client-batch-size", "8",
+		"--auto-truncate",
+	},
+	Env: []corev1.EnvVar{
+		{Name: "MODEL_NAME", Value: "$(spec.modelName)"},
+		{Name: "PORT", Value: "$(spec.runtime.port)"},
+	},
+	VolumeMounts: []corev1.VolumeMount{
+		{Name: "model", MountPath: TEI_MODEL_PATH},
+	},
+	Port: 80,
+}

+ 4 - 0
pkg/templates/templates.go

@@ -28,6 +28,7 @@ import (
 	"github.com/LocoStack/loco-operator/pkg/templates/phoenix"
 	"github.com/LocoStack/loco-operator/pkg/templates/postgresql"
 	"github.com/LocoStack/loco-operator/pkg/templates/qdrant"
+	"github.com/LocoStack/loco-operator/pkg/templates/tei"
 	corev1 "k8s.io/api/core/v1"
 )
 
@@ -257,4 +258,7 @@ func init() {
 	Manager.RegisterRuntime(llamacpp.LlamaCPPTemplate.Spec.Runtime)
 	Manager.RegisterRuntime(llamacpp.LlamaCPPEmbeddingRuntime)
 	Manager.RegisterRuntime(llamacpp.LlamaCPPRerankerRuntime)
+	Manager.RegisterTemplate(tei.TEITemplate)
+	Manager.RegisterRuntime(tei.TEITemplate.Spec.Runtime)
+	Manager.RegisterRuntime(tei.TEIARM64Runtime)
 }