deploy

Requirements.txt

@@ -6,4 +6,5 @@ scikit-learn
 numpy
 pandas
 openpyxl
-numpy
+numpy
+huggingface-hub

service/prediction_service.py

@@ -1,12 +1,34 @@
 import pickle
 from sentence_transformers import SentenceTransformer
-import numpy as np
+import os
+from huggingface_hub import hf_hub_download
 
-# Load Model and Transformer
-with open('models/logistic_regression_model.pkl', 'rb') as f:
-    logistic_model = pickle.load(f)
+# Get the Hugging Face token from environment variable
+hf_token = os.getenv('HF_TOKEN')
+
+# Hugging Face Model ID and local model directory
+hf_model_id = 'Alibaba-NLP/gte-base-en-v1.5'
+model_dir = 'models/sentence_transformer'
+
+# Create model directory if not exists
+os.makedirs(model_dir, exist_ok=True)
+
+# Download model if not already downloaded
+if not os.path.exists(os.path.join(model_dir, 'config.json')):
+    print(f"Downloading model '{hf_model_id}' from Hugging Face...")
+    model_path = hf_hub_download(
+        repo_id=hf_model_id,
+        filename='config.json',
+        cache_dir=model_dir,
+        token=hf_token
+    )
+    # Load model from Hugging Face with token
+    model = SentenceTransformer(hf_model_id, use_auth_token=hf_token, trust_remote_code=True)
+    model.save(model_dir)
+else:
+    print(f"Loading model from local directory: {model_dir}")
+    model = SentenceTransformer(model_dir)
 
-model = SentenceTransformer('models/sentence_transformer')
 
 def predict_label(message: str):
     embedding = model.encode([message])