From 4b8c5aa3264e2633e0b9d288b187bc7ea2abd1d1 Mon Sep 17 00:00:00 2001
From: helloyongyang <yongyang1030@163.com>
Date: Wed, 7 Aug 2024 15:15:22 +0800
Subject: [PATCH] add print_important_package_version log

---
 llmc/__main__.py       | 5 ++++-
 llmc/utils/__init__.py | 3 ++-
 llmc/utils/utils.py    | 9 +++++++++
 3 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/llmc/__main__.py b/llmc/__main__.py
index 5ae500c8..79b44a22 100644
--- a/llmc/__main__.py
+++ b/llmc/__main__.py
@@ -15,7 +15,8 @@
 from llmc.data import BaseDataset, BaseTokenizer
 from llmc.eval import PerplexityEval
 from llmc.models import *
-from llmc.utils import check_config, mkdirs, seed_all
+from llmc.utils import (check_config, mkdirs, print_important_package_version,
+                        seed_all)
 from llmc.utils.registry_factory import ALGO_REGISTRY, MODEL_REGISTRY
 
 
@@ -121,6 +122,8 @@ def main(config):
     logger.info(f'args: {args}')
     logger.info(f'config:\n{json.dumps(config, ensure_ascii=False, indent=4)}')
 
+    print_important_package_version()
+
     seed_all(config.base.seed)
 
     # mkdirs
diff --git a/llmc/utils/__init__.py b/llmc/utils/__init__.py
index bad8b4d1..849c8cd9 100644
--- a/llmc/utils/__init__.py
+++ b/llmc/utils/__init__.py
@@ -1 +1,2 @@
-from .utils import check_config, copy_files, mkdirs, seed_all
+from .utils import (check_config, copy_files, mkdirs,
+                    print_important_package_version, seed_all)
diff --git a/llmc/utils/utils.py b/llmc/utils/utils.py
index 4c198b1c..75a8972d 100644
--- a/llmc/utils/utils.py
+++ b/llmc/utils/utils.py
@@ -72,3 +72,12 @@ def copy_files(source_dir, target_dir, substring):
             target_file = os.path.join(target_dir, filename)
             shutil.copy(source_file, target_file)
             logger.info(f'Copied {filename} to {target_dir}')
+
+
+def print_important_package_version():
+    from importlib.metadata import version
+    logger.info(f"torch : {version('torch')}")
+    logger.info(f"transformers : {version('transformers')}")
+    logger.info(f"tokenizers : {version('tokenizers')}")
+    logger.info(f"huggingface-hub : {version('huggingface-hub')}")
+    logger.info(f"datasets : {version('datasets')}")