update normalize method

This reverts commit 2bf403dca7.

.dockerignore

@@ -1,2 +1,3 @@
 __pycache__
-/venv
+/venv
+/huggingface

.gitattributes

@@ -0,0 +1 @@
+models/** filter=lfs diff=lfs merge=lfs -text

.gitlab-ci.yml

@@ -0,0 +1,36 @@
+# To contribute improvements to CI/CD templates, please follow the Development guide at:
+# https://docs.gitlab.com/ee/development/cicd/templates.html
+# This specific template is located at:
+# https://gitlab.com/gitlab-org/gitlab/-/blob/master/lib/gitlab/ci/templates/Docker.gitlab-ci.yml
+
+# Build a Docker image with CI/CD and push to the GitLab registry.
+# Docker-in-Docker documentation: https://docs.gitlab.com/ee/ci/docker/using_docker_build.html
+#
+# This template uses one generic job with conditional builds
+# for the default branch and all other (MR) branches.
+
+deploy:
+  # Use the official docker image.
+  image: docker:cli
+  stage: deploy
+  services:
+    - docker:dind
+  variables:
+    CI_REGISTRY: registry.waykey.net:7999
+    CI_REGISTRY_IMAGE: $CI_REGISTRY/spiderman/datamining/local-embedding-api
+    DOCKER_IMAGE_NAME: $CI_REGISTRY_IMAGE:$CI_COMMIT_REF_NAME
+  before_script:
+    - chmod 600 $CI_SSH_PRIVATE_KEY
+    - docker login -u "$CI_REGISTRY_USER" -p "$CI_REGISTRY_PASSWORD" $CI_REGISTRY
+  script:
+    - docker build --pull -t "$DOCKER_IMAGE_NAME" .
+    - docker push "$DOCKER_IMAGE_NAME"
+    - |
+      if [[ "$CI_COMMIT_BRANCH" == "$CI_DEFAULT_BRANCH" ]]; then
+        docker tag "$DOCKER_IMAGE_NAME" "$CI_REGISTRY_IMAGE:latest"
+        docker push "$CI_REGISTRY_IMAGE:latest"
+      fi
+  rules:
+    - if: $CI_COMMIT_BRANCH
+      exists:
+        - Dockerfile

Dockerfile

@@ -10,6 +10,4 @@ RUN pip install -r requirements_version.txt --no-cache-dir -i https://pypi.tuna.
 
 COPY . .
 
-RUN python acge_embedding.py
-
 CMD ["python", "main.py"]

acge_embedding.py

@@ -1,17 +1,18 @@
 from transformers import AutoModel, AutoTokenizer
 from sklearn.preprocessing import normalize
 import torch
+import torch.nn.functional as F
 
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print("Using device:", device)
 
-model_name = "aspire/acge-large-zh"
+model_name = "models/aspire--acge-large-zh"
 print("Loading model", model_name)
 model = (
-    AutoModel.from_pretrained(model_name, torch_dtype=torch.float16).eval().to(device)
+    AutoModel.from_pretrained(model_name, local_files_only=True).eval().to(device)
 )
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=True)
 print("Model", model_name, "loaded!")
 
 
@@ -43,9 +44,7 @@ def acge_embedding(text: list[str]) -> list[list[float]]:
         ~attention_mask[..., None].bool(), 0.0
     )
     vector = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
-    vector = normalize(
-        vector.cpu().detach().numpy(),
-        norm="l2",
-        axis=1,
-    )
-    return vector.tolist()
+    # Normalize the output vectors
+    normalized_vector = F.normalize(vector, p=2, dim=1)
+    return normalized_vector.tolist()
+

local-embedding-api.yaml

@@ -0,0 +1,62 @@
+apiVersion: apps/v1
+kind: Deployment
+metadata:
+  name: local-embedding-api
+  labels:
+    app: local-embedding-api
+spec:
+  replicas: 2
+  selector:
+    matchLabels:
+      app: local-embedding-api
+  template:
+    metadata:
+      labels:
+        app: local-embedding-api
+    spec:
+      imagePullSecrets:
+      - name: spiderman-regcred
+      containers:
+      - name: local-embedding-api
+        image: registry.waykey.net:7999/spiderman/datamining/local-embedding-api:latest
+        imagePullPolicy: Always
+        ports:
+        - containerPort: 7999
+        resources:
+          requests:
+            memory: "2000M"
+          limits:
+            memory: "4000M"
+---
+apiVersion: v1
+kind: Service
+metadata:
+  labels:
+    app: local-embedding-api
+  name: local-embedding-api
+  namespace: default
+spec:
+  ports:
+  - port: 7999
+    protocol: TCP
+    targetPort: 7999
+  selector:
+    app: local-embedding-api
+---
+apiVersion: networking.k8s.io/v1
+kind: Ingress
+metadata:
+  name: local-embedding-api
+spec:
+  ingressClassName: nginx
+  rules:
+  - host: local-embedding-api.k8s.local
+    http:
+      paths:
+      - path: /
+        pathType: Prefix
+        backend:
+          service:
+            name: local-embedding-api
+            port:
+              number: 7999