Merge pull request #124 from kerthcet/cleanup/add-log

Add verbose log to modelLoader

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"
     )

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}")

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}")

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")

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"
 )