From 1accdcd8a27ab77c826d610f064c825ba0e8c5c5 Mon Sep 17 00:00:00 2001
From: kerthcet <kerthcet@gmail.com>
Date: Wed, 4 Sep 2024 19:21:00 +0800
Subject: [PATCH] Add verbose log to modelLoader

Signed-off-by: kerthcet <kerthcet@gmail.com>
---
 llmaz/main.py                               | 10 ++-------
 llmaz/model_loader/model_hub/huggingface.py |  7 +++++--
 llmaz/model_loader/model_hub/modelscope.py  |  7 +++++--
 llmaz/util/__init__.py                      |  0
 llmaz/util/logger.py                        | 23 +++++++++++++++++++++
 pkg/defaults.go                             |  2 +-
 6 files changed, 36 insertions(+), 13 deletions(-)
 create mode 100644 llmaz/util/__init__.py
 create mode 100644 llmaz/util/logger.py

diff --git a/llmaz/main.py b/llmaz/main.py
index fd21834..e709874 100644
--- a/llmaz/main.py
+++ b/llmaz/main.py
@@ -15,18 +15,12 @@
 """
 
 import os
-import logging
 from datetime import datetime
 
 from llmaz.model_loader.objstore.objstore import model_download
 from llmaz.model_loader.model_hub.hub_factory import HubFactory
 from llmaz.model_loader.model_hub.huggingface import HUGGING_FACE
-
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-)
-logger = logging.getLogger(__name__)
+from llmaz.util.logger import Logger
 
 
 if __name__ == "__main__":
@@ -54,6 +48,6 @@
     else:
         raise EnvironmentError(f"unknown model source type {model_source_type}")
 
-    logger.info(
+    Logger.info(
         f"loading models from {model_source_type} takes {datetime.now() - start_time}s"
     )
diff --git a/llmaz/model_loader/model_hub/huggingface.py b/llmaz/model_loader/model_hub/huggingface.py
index 0de086a..678957a 100644
--- a/llmaz/model_loader/model_hub/huggingface.py
+++ b/llmaz/model_loader/model_hub/huggingface.py
@@ -25,6 +25,7 @@
     MAX_WORKERS,
     ModelHub,
 )
+from llmaz.util.logger import Logger
 
 from typing import Optional
 
@@ -38,7 +39,9 @@ def name(cls) -> str:
     def load_model(
         cls, model_id: str, filename: Optional[str], revision: Optional[str]
     ) -> None:
-        print(f"Start to download model {model_id}")
+        Logger.info(
+            f"Start to download, model_id: {model_id}, filename: {filename}, revision: {revision}"
+        )
 
         if filename:
             hf_hub_download(
@@ -71,4 +74,4 @@ def load_model(
 
 def handle_completion(future):
     filename = future.result()
-    print(f"Download completed for {filename}")
+    Logger.info(f"Download completed for {filename}")
diff --git a/llmaz/model_loader/model_hub/modelscope.py b/llmaz/model_loader/model_hub/modelscope.py
index eb05b7b..d7eaeac 100644
--- a/llmaz/model_loader/model_hub/modelscope.py
+++ b/llmaz/model_loader/model_hub/modelscope.py
@@ -26,6 +26,7 @@
     MODEL_SCOPE,
     ModelHub,
 )
+from llmaz.util.logger import Logger
 
 
 class ModelScope(ModelHub):
@@ -38,7 +39,9 @@ def name(cls) -> str:
     def load_model(
         cls, model_id: str, filename: Optional[str], revision: Optional[str]
     ) -> None:
-        print(f"Start to download model {model_id}")
+        Logger.info(
+            f"Start to download, model_id: {model_id}, filename: {filename}, revision: {revision}"
+        )
 
         with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
             futures = []
@@ -57,4 +60,4 @@ def load_model(
 
 def handle_completion(future):
     filename = future.result()
-    print(f"Download completed for {filename}")
+    Logger.info(f"Download completed for {filename}")
diff --git a/llmaz/util/__init__.py b/llmaz/util/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/llmaz/util/logger.py b/llmaz/util/logger.py
new file mode 100644
index 0000000..a2b01ee
--- /dev/null
+++ b/llmaz/util/logger.py
@@ -0,0 +1,23 @@
+"""
+Copyright 2024.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import logging
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+Logger = logging.getLogger("llmaz")
diff --git a/pkg/defaults.go b/pkg/defaults.go
index 9fdc95f..f854aef 100644
--- a/pkg/defaults.go
+++ b/pkg/defaults.go
@@ -17,5 +17,5 @@ limitations under the License.
 package pkg
 
 const (
-	LOADER_IMAGE = "inftyai/model-loader:v0.0.6"
+	LOADER_IMAGE = "inftyai/model-loader:v0.0.7"
 )