Merge pull request #15 from Sistem-Cerdas-Recruitment/production

cloudbuild.yaml

@@ -1,16 +1,12 @@
 steps:
   - name: "gcr.io/cloud-builders/docker"
     dir: "core-model-prediction"
-    args:
-      [
-        "build",
-        "--build-arg",
-        "HF_TOKEN=${_HF_TOKEN}",
-        "-t",
-        "us-central1-docker.pkg.dev/${PROJECT_ID}/interview-ai-detector/model-prediction:latest",
-        ".",
-      ]
-    secretEnv: ["HF_TOKEN"]
+    entrypoint: "bash"
+    secretEnv: ["_HF_TOKEN"]
+    args: 
+      - -c
+      - | 
+          docker build -t us-central1-docker.pkg.dev/${PROJECT_ID}/interview-ai-detector/model-prediction:latest --build-arg HF_TOKEN=$$_HF_TOKEN .
 
   - name: "gcr.io/cloud-builders/docker"
     args:
@@ -19,10 +15,14 @@ steps:
         "us-central1-docker.pkg.dev/${PROJECT_ID}/interview-ai-detector/model-prediction:latest",
       ]
 
+options:
+  pool:
+    name: 'projects/${PROJECT_ID}/locations/us-central1/workerPools/ai-detector-builder'
+
 images:
   - "us-central1-docker.pkg.dev/${PROJECT_ID}/interview-ai-detector/model-prediction:latest"
 
 availableSecrets:
   secretManager:
-    - versionName: "projects/${PROJECT_ID}/secrets/HF_TOKEN/versions/1"
-      env: "HF_TOKEN"
+    - versionName: "projects/${PROJECT_ID}/secrets/HF_TOKEN/versions/latest"
+      env: "_HF_TOKEN"

core-model-prediction/Dockerfile

@@ -1,6 +1,3 @@
-# HF Token args
-ARG HF_TOKEN
-
 # Use an official Python runtime as a base image
 FROM pytorch/pytorch:2.1.2-cuda12.1-cudnn8-runtime
 
@@ -20,6 +17,9 @@ RUN python -m nltk.downloader punkt wordnet averaged_perceptron_tagger
 # Unzip wordnet
 RUN unzip /root/nltk_data/corpora/wordnet.zip -d /root/nltk_data/corpora/
 
+# HF Token args
+ARG HF_TOKEN
+
 # Download HuggingFace model
 RUN python -c "from transformers import AutoTokenizer, AutoModelForCausalLM; \
     tokenizer = AutoTokenizer.from_pretrained('google/gemma-2b', token='$HF_TOKEN'); \
@@ -28,7 +28,7 @@ RUN python -c "from transformers import AutoTokenizer, AutoModelForCausalLM; \
     model.save_pretrained('/app/gemma-2b')"
 
 # Model env
-ENV MODEL_DIR=/app/gemma-2b
+ENV MODEL_DIR=gemma-2b
 
 # Make port 8080 available to the world outside this container
 EXPOSE 8080

core-model-prediction/gemma2b_dependencies.py

@@ -13,9 +13,9 @@ class Gemma2BDependencies:
     def __new__(cls):
         if cls._instance is None:
             cls._instance = super(Gemma2BDependencies, cls).__new__(cls)
-            model_dir = os.getenv("MODEL_DIR", "/app/gemma-2b")
-            cls._instance.tokenizer = AutoTokenizer.from_pretrained(model_dir)
-            cls._instance.model = AutoModelForCausalLM.from_pretrained(model_dir)
+            model_dir = os.getenv("MODEL_DIR", "gemma-2b")
+            cls._instance.tokenizer = AutoTokenizer.from_pretrained(model_dir, local_files_only=True)
+            cls._instance.model = AutoModelForCausalLM.from_pretrained(model_dir, local_files_only=True)
             cls._instance.device = DeviceManager()
             cls._instance.model.to(cls._instance.device)
         return cls._instance