# /home/user/app/utils/download_model.py
import os
from transformers import AutoModel

def download_model():
    # Use environment variables for the model name and directory
    model_name = os.getenv("HF_MODEL_NAME")
    model_dir = f"/home/user/data/models/{model_name}"

    # Authenticate with Hugging Face using the token, if available
    hf_token = os.getenv("HF_TOKEN")
    if hf_token:
        from huggingface_hub import HfFolder
        HfFolder.save_token(hf_token)  # Save the token for later use by the library

    # Download the model
    print(f"Downloading model: {model_name}...")
    model = AutoModel.from_pretrained(model_name)
    model.save_pretrained(model_dir)
    print(f"Model {model_name} downloaded and saved to {model_dir}")

if __name__ == "__main__":
    download_model()