From 46e60db4f1d607c2aec5647c939b568ebfcb381b Mon Sep 17 00:00:00 2001
From: Michal Januszewski <mjanusz@google.com>
Date: Sun, 11 Aug 2024 06:19:01 -0700
Subject: [PATCH] Add training utilities.

PiperOrigin-RevId: 661798999
---
 ffn/training/import_util.py | 34 +++++-----------------------------
 1 file changed, 5 insertions(+), 29 deletions(-)

diff --git a/ffn/training/import_util.py b/ffn/training/import_util.py
index 61c2dc6..6dc8a4b 100644
--- a/ffn/training/import_util.py
+++ b/ffn/training/import_util.py
@@ -14,34 +14,10 @@
 # ==============================================================================
 """Functions for dynamically importing symbols from modules."""
 
-import importlib
-from absl import logging
+from connectomics.common import import_util
 
 
-def import_symbol(specifier, default_packages='ffn.training.models'):
-  """Imports a symbol from a python module.
-
-  The calling module must have the target module for the import as dependency.
-
-  Args:
-    specifier: full path specifier in format
-        [<packages>.]<module_name>.<model_class>, if packages is missing
-        ``default_packages`` is used.
-    default_packages: chain of packages before module in format
-        <top_pack>.<sub_pack>.<subsub_pack> etc.
-
-  Returns:
-    symbol: object from module
-  """
-  module_path, symbol_name = specifier.rsplit('.', 1)
-  try:
-    logging.info('Importing symbol %s from %s.%s',
-                 symbol_name, default_packages, module_path)
-    module = importlib.import_module(default_packages + '.' + module_path)
-  except ImportError as e:
-    logging.info(e)
-    logging.info('Importing symbol %s from %s', symbol_name, module_path)
-    module = importlib.import_module(module_path)
-
-  symbol = getattr(module, symbol_name)
-  return symbol
+def import_symbol(
+    specifier: str, default_packages: str = 'ffn.training.models'
+):
+  return import_util.import_symbol(specifier, default_packages)