diff --git a/README.md b/README.md index 57c9027..01a6069 100644 --- a/README.md +++ b/README.md @@ -31,7 +31,7 @@ pip install modelscan With it installed, scan a model: ```bash -modelscan -p /path/to/model_file.h5 +modelscan -p /path/to/model_file.pkl ``` ## Why You Should Scan Models @@ -105,6 +105,11 @@ or `pyproject.toml` like this: modelscan = ">=0.1.1" ``` +Scanners for Tensorflow or HD5 formatted models require installation with extras: +```bash +pip install 'modelscan[ tensorflow, h5py ]' +``` + ### Using ModelScan via CLI ModelScan supports the following arguments via the CLI: diff --git a/modelscan/scanners/saved_model/scan.py b/modelscan/scanners/saved_model/scan.py index 357b047..74e8fb8 100644 --- a/modelscan/scanners/saved_model/scan.py +++ b/modelscan/scanners/saved_model/scan.py @@ -64,7 +64,9 @@ def _check_for_unsafe_tf_keras_operator( unsafe_operators: Dict[str, Any], ) -> ScanResults: issues: List[Issue] = [] - all_operators = tensorflow.raw_ops.__dict__.keys() + all_operators = ( + tensorflow.raw_ops.__dict__.keys() if tensorflow_installed else [] + ) all_safe_operators = [ operator for operator in list(all_operators) if operator[0] != "_" ]