YorubaCNN / model_loader.py
Testys's picture
Readding documents necessary for Streamlit Document
73b342f
from huggingface_hub import HfApi, HfFolder
import os
import os
from dotenv import load_dotenv
# Load environment variables from .env file if it exists
load_dotenv()
hf_token = os.getenv('HF_TOKEN')
api = HfApi()
# Specify the model details
model_id = "Testys/cnn_yor_ner"
sent_id = "Testys/cnn_sent_yor"
local_dir = "./my_model"
sent_dir = "./sent_model"
# Download the model folder
if not os.path.exists(local_dir):
os.makedirs(local_dir)
if not os.path.exists(os.path.join(local_dir, "pytorch_model.bin")):
api.hf_hub_download(repo_id=model_id, filename="pytorch_model.bin", local_dir=local_dir, use_auth_token=hf_token)
if not os.path.exists(os.path.join(local_dir, "config.json")):
api.hf_hub_download(repo_id=model_id, filename="config.json", local_dir=local_dir, use_auth_token=hf_token)
# Check if the model is already downloaded
if not os.path.exists(sent_dir):
os.makedirs(sent_dir)
# Download the model files only if they don't exist
if not os.path.exists(os.path.join(sent_dir, "sent_pytorch_model.bin")):
api.hf_hub_download(repo_id=sent_id, filename="sent_pytorch_model.bin", local_dir=sent_dir, use_auth_token=hf_token)
if not os.path.exists(os.path.join(sent_dir, "config.json")):
api.hf_hub_download(repo_id=sent_id, filename="config.json", local_dir=sent_dir, use_auth_token=hf_token)