implemented load_model function

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()
+    model: VisualCommunicationModel = load_model(client=minio_client)
-    # 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)
     # show model weights
     summary(model)
     print('loaded model')