From e62435cf22fe229cb8c386836b8b6d70330202d1 Mon Sep 17 00:00:00 2001
From: Brian Bjarke Jensen <bbj@skov.dk>
Date: Mon, 22 Jul 2024 23:47:39 +0200
Subject: [PATCH] added full model load script

---
 model/src/main.py | 53 ++++++++++++++++++++++-------------------------
 1 file changed, 25 insertions(+), 28 deletions(-)

diff --git a/model/src/main.py b/model/src/main.py
index c8fc5d5..888131f 100644
--- a/model/src/main.py
+++ b/model/src/main.py
@@ -1,45 +1,42 @@
 """Main script to be run by service."""
-from __future__ import annotations
 
-from hashlib import md5
-from io import BytesIO
+import logging
+from pathlib import Path
 
 import torch
-from data_store import connect
-from data_store import put
+from dotenv import load_dotenv
 from models import VisualCommunicationModel
 from torchinfo import summary
-from utils import check_env
+from utils import get_model_name
 
+from shared.data_store import connect, get_model
 from shared.utils import setup_logging
 
-
 DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 if __name__ == '__main__':
-    # check that environment variables are set
-    check_env()
+    # 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().to(DEVICE)
+    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)
     # show model weights
     summary(model)
-    # for param in model.parameters():
-    #     logging.info(param.data)
-    # save model to buffer
-    buffer = BytesIO()
-    torch.save(model.state_dict(), buffer)
-    print(f"buffer size: {len(buffer.getvalue())}")
-    # calculate buffer hash
-    hash_str = md5(buffer.getbuffer()).hexdigest()
-    print(f"hash string: {hash_str}")
-    # connect to minio
-    client = connect()
-    # put buffer in minio bucket
-    put(
-        client=client,
-        buffer=buffer,
-        object_name=hash_str,
-    )
-    print('saved to Minio')
+    print('loaded model')