fix main README examples and add option to limit search response time (…

…neuralmagic#197)

* fix main README examples and add option to limit search response time

* use search return_stubs in tests and examples

.gitignore

@@ -1,6 +1,9 @@
 # Created by https://www.toptal.com/developers/gitignore/api/windows,macos,linux,python,react,pycharm,emacs,vim,visualstudio,visualstudiocode
 # Edit at https://www.toptal.com/developers/gitignore?templates=windows,macos,linux,python,react,pycharm,emacs,vim,visualstudio,visualstudiocode
 
+# notebook results
+notebooks/model_repo/
+
 # envs
 venv*/
 

README.md

@@ -86,37 +86,32 @@ pip install sparsezoo
 ### Python APIs
 
 The Python APIs respect this format enabling you to search and download models. Some code examples are given below.
+The [SparseZoo UI](https://sparsezoo.neuralmagic.com/) also enables users to load models by copying
+a stub directly from a model page.
 
-#### Searching the Zoo
-
-```python
-from sparsezoo import Zoo
-
-models = Zoo.search_models(domain="cv", sub_domain="classification")
-print(models)
-```
 
-#### Common Models
+#### Loading from a Stub
 
 ```python
-from sparsezoo.models.classification import resnet_50
-
-model = resnet_50()
-model.download()
+from sparsezoo import Model
 
-print(model.onnx_file.downloaded_path())
+# copied from https://sparsezoo.neuralmagic.com/
+stub = "zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenet/pruned90_quant-none"
+model = Model(stub)
+print(model)
 ```
 
-#### Searching Optimized Versions
+#### Searching the Zoo
 
 ```python
-from sparsezoo import Zoo
-from sparsezoo.models.classification import resnet_50
+from sparsezoo import search_models
 
-search_model = resnet_50()
-sparse_models = Zoo.search_sparse_models(search_model)
-
-print(sparse_models)
+models = search_models(
+    domain="cv",
+    sub_domain="classification",
+    return_stubs=True,
+)
+print(models)
 ```
 
 ### Environmental Variables
@@ -158,10 +153,10 @@ Search command help
 sparsezoo search -h
 ```
 
-<br>Searching for all classification models in the computer vision domain
+<br>Searching for all classification MobileNetV1 models in the computer vision domain
 
 ```shell script
-sparsezoo search --domain cv --sub-domain classification
+sparsezoo search --domain cv --sub-domain classification --architecture mobilenet_v1
 ```
 
 <br>Searching for all ResNet-50 models

notebooks/model_download.ipynb

@@ -54,6 +54,7 @@
     "try:\n",
     "    # make sure sparsezoo is installed\n",
     "    import sparsezoo\n",
+    "    print(\"sparsezoo available for model download\")\n",
     "except Exception as ex:\n",
     "    raise Exception(\n",
     "        \"please install sparsezoo using the setup.py file before continuing\"\n",
@@ -99,16 +100,16 @@
    "source": [
     "## Selecting a model using the search functionality ##\n",
     "\n",
-    "n_first_print = 5 # print out first n models found\n",
-    "domain, sub_domain = \"cv\", \"classification\" # defining the search criteria\n",
+    "n_first_print = 5  # print out first n models found\n",
+    "domain, sub_domain = \"cv\", \"classification\"  # defining the search criteria\n",
     "\n",
     "print(\"Searching for models...\")\n",
-    "models = search_models(domain, sub_domain)\n",
-    "print(f\"Given the search criteria, found {len(models)} models...\")\n",
-    "print(f\"Printing first {n_first_print} models...\")\n",
-    "[print(f\"\\t{model}\") for model in models[:n_first_print]]\n",
-    "model = models[0]\n",
-    "print(f\"\\nSelecting the first model: {str(model)}.\")"
+    "model_stubs = search_models(domain, sub_domain, return_stubs=True)\n",
+    "print(f\"Given the search criteria, found {len(model_stubs)} models...\")\n",
+    "print(f\"Printing first {n_first_print} model stubs...\")\n",
+    "[print(f\"\\t{stub}\") for stub in model_stubs[:n_first_print]]\n",
+    "model_stub = model_stubs[0]\n",
+    "print(f\"\\nSelecting the first model: {str(model_stub)}.\")"
    ]
   },
   {
@@ -117,7 +118,7 @@
    "source": [
     "### Step 3 - Downloading the Model\n",
     "\n",
-    "After making a model selection, run the cell block below to download the model locally. By default, it will save the model to an appropriately named folder under the current working directory. You can change the save_dir if you'd like to save to another location on the local system.\n"
+    "After making a model selection, run the cell block below to download the model locally from the selected stub. By default, it will save the model to an appropriately named folder under the current working directory. You can change the save_dir if you'd like to save to another location on the local system.\n"
    ]
   },
   {
@@ -127,7 +128,8 @@
    "outputs": [],
    "source": [
     "import os\n",
-    "model.download(directory_path = download_path, override = True)\n",
+    "model = Model(model_stub)\n",
+    "# calling .path triggers model download\n",
     "print(f\"Model download to {model.path}\")"
    ]
   },
@@ -159,7 +161,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.8.0"
   }
  },
  "nbformat": 4,

src/sparsezoo/search/search.py

@@ -38,6 +38,7 @@ def search_models(
     page: int = 1,
     page_length: int = 20,
     force_token_refresh: bool = False,
+    return_stubs: bool = False,
 ) -> List[Model]:
     """
     Obtains a list of Models matching the search parameters
@@ -72,6 +73,8 @@ def search_models(
     :param override_parent_path: Path to override the default save path
         for where to save the parent folder for this file under
     :param force_token_refresh: True to refresh the auth token, False otherwise
+    :param return_stubs: if True, found models will be returned as stubs only
+        instead of retrieving full info for each found model. Default False
     :return: The requested list of Model instances
     """
     args = {
@@ -99,10 +102,12 @@ def search_models(
         force_token_refresh=force_token_refresh,
     )
 
-    return [
-        Model(model_args_to_stub(**model_dict))
-        for model_dict in response_json["models"]
-    ]
+    stubs = [model_args_to_stub(**model_dict) for model_dict in response_json["models"]]
+
+    if return_stubs:
+        return stubs
+
+    return [Model(stub) for stub in stubs]
 
 
 def model_args_to_stub(**kwargs) -> str:

tests/sparsezoo/search/test_search.py

@@ -37,5 +37,5 @@
     scope="function",
 )
 def test_search_models(search_args):
-    models = search_models(**search_args)
+    models = search_models(**search_args, return_stubs=True)
     assert len(models) > 0