diff --git a/llmaz/main.py b/llmaz/main.py index fd21834f..e709874b 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 0de086aa..678957a3 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 eb05b7b4..d7eaeac9 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 00000000..e69de29b diff --git a/llmaz/util/logger.py b/llmaz/util/logger.py new file mode 100644 index 00000000..a2b01ee7 --- /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 9fdc95fb..f854aefb 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" )