EduConnect / app /utils /download_model.py
dtyago's picture
Entrypoint to load model
faf4679
raw
history blame
806 Bytes
# /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()