From 89aecfdef7076180c0b3e7b18c8a3fe89af5dbe0 Mon Sep 17 00:00:00 2001
From: Brian Bjarke Jensen <bbj@skov.dk>
Date: Sun, 14 Jul 2024 11:16:35 +0200
Subject: [PATCH] updated to use new io functions

---
 tests/model_io_test.py | 27 ++++++++++++---------------
 1 file changed, 12 insertions(+), 15 deletions(-)

diff --git a/tests/model_io_test.py b/tests/model_io_test.py
index 9eccc96..bde5a0b 100644
--- a/tests/model_io_test.py
+++ b/tests/model_io_test.py
@@ -1,37 +1,34 @@
 """Test that a model can be saved and loaded again."""
 
-from io import BytesIO
+from pathlib import Path
 
 import torch
+from dotenv import load_dotenv
 from torchinfo import summary
 
 from model.src.models import VisualCommunicationModel
-from shared.data_store import connect, put
-from shared.utils import check_env, setup_logging
+from shared.data_store import connect, put_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 / 'server.env'
+    assert env_path.exists()
+    load_dotenv(env_path)
     # setup logging
     setup_logging()
+    # connect to minio
+    client = connect()
     # 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())}")
-    # connect to minio
-    client = connect(bucket='models')
     # put buffer in minio bucket
-    hash_str = put(
+    hash_str = put_model(
         client=client,
-        buffer=buffer,
+        model=model,
     )
     print(f"hash string: {hash_str}")
     print('saved to Minio')