updated main script

model/src/main.py

@@ -0,0 +1,45 @@
+"""Main script to be run by service."""
+from __future__ import annotations
+
+from hashlib import md5
+from io import BytesIO
+
+import torch
+from min_io import connect
+from min_io import put
+from models import VisualCommunicationModel
+from torchinfo import summary
+from utils import check_env
+
+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()
+    # setup logging
+    setup_logging()
+    # instantiate model
+    model = VisualCommunicationModel().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')