diff --git a/model/src/main.py b/model/src/main.py index 888131f..e776beb 100644 --- a/model/src/main.py +++ b/model/src/main.py @@ -1,42 +1,25 @@ """Main script to be run by service.""" -import logging from pathlib import Path -import torch from dotenv import load_dotenv from models import VisualCommunicationModel from torchinfo import summary -from utils import get_model_name +from utils import load_model -from shared.data_store import connect, get_model +from shared.data_store import connect from shared.utils import setup_logging -DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') - if __name__ == '__main__': # load in env file env_path = Path(__file__).parent.parent.parent / 'server.env' - assert env_path.exists() load_dotenv(env_path) # setup logging setup_logging() # connect to minio minio_client = connect() # instantiate model - model = VisualCommunicationModel() - # get model object name - model_name_path = Path(__file__).parent / 'model_name.txt' - model_object_name = get_model_name(path=model_name_path) - logging.info('using model: %s', model_object_name) - # load model from minio - model_checkpoint = get_model( - client=minio_client, - object_name=model_object_name, - ) - model.load_state_dict(model_checkpoint) - # move model to selected device - model = model.to(DEVICE) + model: VisualCommunicationModel = load_model(client=minio_client) # show model weights summary(model) print('loaded model')