replace compare_versions from PL

flash/core/utilities/imports.py

@@ -17,8 +17,9 @@
 import os
 import types
 from importlib.util import find_spec
-from typing import List, Tuple, Union
+from typing import Callable, List, Tuple, Union
 
+import pkg_resources
 from pkg_resources import DistributionNotFound
 
 try:
@@ -48,26 +49,29 @@ def _module_available(module_path: str) -> bool:
         return True
 
 
-def _compare_version(package: str, op, version) -> bool:
+def _compare_version(package: str, op: Callable, version: str, use_base_version: bool = False) -> bool:
     """Compare package version with some requirements.
 
     >>> _compare_version("torch", operator.ge, "0.1")
     True
+    >>> _compare_version("does_not_exist", operator.ge, "0.0")
+    False
     """
     try:
         pkg = importlib.import_module(package)
-    except (ModuleNotFoundError, DistributionNotFound, ValueError):
+    except (ImportError, DistributionNotFound):
         return False
     try:
-        pkg_version = Version(pkg.__version__)
-    except AttributeError:
-        # In case the module doesn't have __version__ attribute (example: baal)
-        import pkg_resources
-
-        pkg_version = Version(pkg_resources.get_distribution("baal").version)
+        if hasattr(pkg, "__version__"):
+            pkg_version = Version(pkg.__version__)
+        else:
+            # try pkg_resources to infer version
+            pkg_version = Version(pkg_resources.get_distribution(package).version)
     except TypeError:
-        # this is mock by sphinx, so it shall return True to generate all summaries
+        # this is mocked by Sphinx, so it should return True to generate all summaries
         return True
+    if use_base_version:
+        pkg_version = Version(pkg_version.base_version)
     return op(pkg_version, Version(version))