deploy

service/prediction_service.py

@@ -1,4 +1,3 @@
-import pickle
 from sentence_transformers import SentenceTransformer
 import os
 from huggingface_hub import hf_hub_download
@@ -8,7 +7,7 @@ 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'
+model_dir = '/tmp/sentence_transformer'  # Use /tmp for write permissions
 
 # Create model directory if not exists
 os.makedirs(model_dir, exist_ok=True)
@@ -16,23 +15,8 @@ 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)
-
-
-def predict_label(message: str):
-    embedding = model.encode([message])
-    prediction = logistic_model.predict(embedding)[0]
-    probability = logistic_model.predict_proba(embedding)[0].max()
-    
-    return prediction, float(probability)