File size: 837 Bytes
93109de 1eb51e0 93109de 1eb51e0 93109de 1eb51e0 93109de 1eb51e0 93109de 1eb51e0 93109de |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 |
"""
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
|