from sentence_transformers import SentenceTransformer
import os
from huggingface_hub import login

# Get the token from environment variable
hf_token = os.getenv('HF_TOKEN')

# Model path and Hugging Face model ID
model_dir = '/tmp/sentence_transformer'
hf_model_id = 'Alibaba-NLP/gte-base-en-v1.5'

# Authenticate with Hugging Face
login(token=hf_token)

# Check if model directory exists
if not os.path.exists(model_dir):
    print(f"Downloading model '{hf_model_id}' from Hugging Face...")
    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)