|
| 1 | +# Copyright 2023 The TensorFlow Authors. All Rights Reserved. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | +# ============================================================================== |
| 15 | +"""Package dependencies for model-card-toolkit.""" |
| 16 | + |
| 17 | +import importlib |
| 18 | +from typing import Dict, List |
| 19 | + |
| 20 | +_VERSIONS = { |
| 21 | + 'absl': 'absl-py>=0.9,<1.1', |
| 22 | + 'importlib_resources': 'importlib-resources>=1.3.0; python_version<"3.9"', |
| 23 | + 'isort': 'isort', |
| 24 | + 'jinja2': 'jinja2>=3.1,<3.2', |
| 25 | + 'jsonschema': 'jsonschema>=3.2.0,<4', |
| 26 | + 'matplotlib': 'matplotlib>=3.2.0,<4', |
| 27 | + 'ml_metadata': 'ml-metadata>=1.5.0,<2.0.0', |
| 28 | + 'pre-commit': 'pre-commit', |
| 29 | + 'protobuf': 'protobuf>=3.19.0,<4', |
| 30 | + 'pylint': 'pylint', |
| 31 | + 'pytest': 'pytest', |
| 32 | + 'tensorflow_data_validation': 'tensorflow-data-validation>=1.5.0,<2.0.0', |
| 33 | + 'tensorflow_datasets': 'tensorflow-datasets>=4.8.2', |
| 34 | + 'tensorflow_metadata': 'tensorflow-metadata>=1.5.0,<2.0.0', |
| 35 | + 'tensorflow_model_analysis': 'tensorflow-model-analysis>=0.36.0,<0.45.0', |
| 36 | + 'yapf': 'yapf', |
| 37 | +} |
| 38 | + |
| 39 | +_REQUIRED_DEPS = [ |
| 40 | + 'importlib_resources', # reading resource files, e.g. model card templates |
| 41 | + 'jinja2', # rendering model card templates |
| 42 | + 'jsonschema', # validating JSON schema |
| 43 | + 'matplotlib', # plotting |
| 44 | + 'protobuf', # working with model card protos |
| 45 | +] |
| 46 | + |
| 47 | +_EXAMPLES_EXTRA_DEPS = [ |
| 48 | + # Required for model_card_toolkit.documentation.examples.cats_vs_dogs |
| 49 | + 'tensorflow_datasets', |
| 50 | +] |
| 51 | + |
| 52 | +_TENSORFLOW_EXTRA_DEPS = [ |
| 53 | + 'ml_metadata', |
| 54 | + 'tensorflow_data_validation', |
| 55 | + 'tensorflow_metadata', |
| 56 | + 'tensorflow_model_analysis', |
| 57 | +] |
| 58 | + |
| 59 | +_TEST_EXTRA_DEPS = ['absl', 'isort', 'pre-commit', 'pylint', 'pytest', 'yapf'] |
| 60 | + |
| 61 | +TENSORFLOW_EXTRA_IMPORT_ERROR_MSG = """ |
| 62 | +This functionaliy requires `tensorflow` extra dependencies but they were not |
| 63 | +found in your environment. You can install them with: |
| 64 | +``` |
| 65 | +pip install model-card-toolkit[tensorflow] |
| 66 | +``` |
| 67 | +""" |
| 68 | + |
| 69 | + |
| 70 | +def _make_deps_list(package_names: List[str]) -> List[str]: |
| 71 | + """Returns a list of dependencies with their constraints. |
| 72 | +
|
| 73 | + Raises: ValueError if a `package_name` is not in the list of known dependencies. |
| 74 | + """ |
| 75 | + deps = [] |
| 76 | + for package_name in package_names: |
| 77 | + if package_name not in _VERSIONS: |
| 78 | + raise ValueError( |
| 79 | + f'Package {package_name} is not in the list of known dependencies: ' |
| 80 | + f'{_VERSIONS.keys()}' |
| 81 | + ) |
| 82 | + deps.append(_VERSIONS[package_name]) |
| 83 | + return deps |
| 84 | + |
| 85 | + |
| 86 | +def make_required_install_packages() -> List[str]: |
| 87 | + """Returns the list of required packages.""" |
| 88 | + return _make_deps_list(_REQUIRED_DEPS) |
| 89 | + |
| 90 | + |
| 91 | +def make_extra_packages_examples() -> List[str]: |
| 92 | + """Returns the list of packages needed for running examples.""" |
| 93 | + return _make_deps_list(_EXAMPLES_EXTRA_DEPS) |
| 94 | + |
| 95 | + |
| 96 | +def make_extra_packages_tensorflow() -> List[str]: |
| 97 | + """Returns the list of packages needed to use TensorFlow utils.""" |
| 98 | + return _make_deps_list(_TENSORFLOW_EXTRA_DEPS) |
| 99 | + |
| 100 | + |
| 101 | +def has_tensorflow_extra_deps() -> bool: |
| 102 | + """Returns True if all tensorflow extra dependencies are installed.""" |
| 103 | + return all(importlib.util.find_spec(name) for name in _TENSORFLOW_EXTRA_DEPS) |
| 104 | + |
| 105 | + |
| 106 | +def ensure_tensorflow_extra_deps_installed(): |
| 107 | + """Raises ImportError if tensorflow extra dependencies are not installed. |
| 108 | + """ |
| 109 | + if not has_tensorflow_extra_deps(): |
| 110 | + raise ImportError(TENSORFLOW_EXTRA_IMPORT_ERROR_MSG) |
| 111 | + |
| 112 | + |
| 113 | +def make_extra_packages_test() -> List[str]: |
| 114 | + """Returns the list of packages needed for running tests.""" |
| 115 | + return _make_deps_list(_TEST_EXTRA_DEPS) |
| 116 | + |
| 117 | + |
| 118 | +def make_extra_packages_all() -> List[str]: |
| 119 | + """Returns the list of all optional packages.""" |
| 120 | + return [ |
| 121 | + *make_extra_packages_examples(), |
| 122 | + *make_extra_packages_tensorflow(), |
| 123 | + *make_extra_packages_test(), |
| 124 | + ] |
| 125 | + |
| 126 | + |
| 127 | +def make_required_extra_packages() -> Dict[str, List[str]]: |
| 128 | + """Returns the dict of required extra packages.""" |
| 129 | + return { |
| 130 | + 'examples': make_extra_packages_examples(), |
| 131 | + 'tensorflow': make_extra_packages_tensorflow(), |
| 132 | + 'test': make_extra_packages_test(), |
| 133 | + 'all': make_extra_packages_all(), |
| 134 | + } |
0 commit comments