"""
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