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