""" | |
Module to load the project models | |
""" | |
import os | |
import tensorflow as tf | |
import tensorflow_hub as hub | |
import tensorflow_text | |
from dotenv import load_dotenv | |
from huggingface_hub import hf_hub_download | |
load_dotenv() | |
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__)) | |
MODEL_FILENAME = os.getenv("MODEL_FILENAME") | |
MODEL_REPOSITORY_NAME = os.getenv("MODEL_REPOSITORY_NAME") | |
def load_sentiments_model(): | |
""" | |
Load pretrained model | |
""" | |
model_path = os.path.join(CURRENT_DIR, MODEL_FILENAME) | |
# If model doesnt exist download from huggingface | |
if not os.path.exists(model_path): | |
hf_hub_download(MODEL_REPOSITORY_NAME, MODEL_FILENAME, local_dir=CURRENT_DIR) | |
model = tf.keras.models.load_model( | |
model_path, custom_objects={"KerasLayer": hub.KerasLayer}, compile=False | |
) | |
return model | |