Search by authors and string (#531)

* token in env variables

* done

* done

* Update src/huggingface_hub/hf_api.py

Co-authored-by: Julien Chaumond <[email protected]>

* search by  and  added in

* fix in test

* Update src/huggingface_hub/hf_api.py

Co-authored-by: Zachary Mueller <[email protected]>

* Update tests/test_hf_api.py

Co-authored-by: Zachary Mueller <[email protected]>

* Update tests/test_hf_api.py

Co-authored-by: Zachary Mueller <[email protected]>

* Update src/huggingface_hub/hf_api.py

Co-authored-by: Zachary Mueller <[email protected]>

* quality + test

Co-authored-by: Julien Chaumond <[email protected]>
Co-authored-by: Zachary Mueller <[email protected]>

Author: 3 people

src/huggingface_hub/hf_api.py

@@ -507,6 +507,8 @@ def get_dataset_tags(self) -> DatasetTags:
     def list_models(
         self,
         filter: Union[str, Iterable[str], None] = None,
+        author: Optional[str] = None,
+        search: Optional[str] = None,
         sort: Union[Literal["lastModified"], str, None] = None,
         direction: Optional[Literal[-1]] = None,
         limit: Optional[int] = None,
@@ -538,6 +540,30 @@ def list_models(
 
                     >>> # List only the models from the AllenNLP library
                     >>> api.list_models(filter="allennlp")
+            author (:obj:`str`, `optional`):
+                A string which identify the author (user or organization) of the returned models
+                Example usage:
+
+                    >>> from huggingface_hub import HfApi
+                    >>> api = HfApi()
+
+                    >>> # List all models from google
+                    >>> api.list_models(author="google")
+
+                    >>> # List only the text classification models from google
+                    >>> api.list_models(filter="text-classification", author="google")
+            search (:obj:`str`, `optional`):
+                A string that will be contained in the returned models
+                Example usage:
+
+                    >>> from huggingface_hub import HfApi
+                    >>> api = HfApi()
+
+                    >>> # List all models with "bert" in their name
+                    >>> api.list_models(search="bert")
+
+                    >>> #List all models with "bert" in their name made by google
+                    >>> api.list_models(search="bert", author="google")
             sort (:obj:`Literal["lastModified"]` or :obj:`str`, `optional`):
                 The key with which to sort the resulting models. Possible values are the properties of the `ModelInfo`
                 class.
@@ -558,6 +584,10 @@ def list_models(
         if filter is not None:
             params.update({"filter": filter})
             params.update({"full": True})
+        if author is not None:
+            params.update({"author": author})
+        if search is not None:
+            params.update({"search": search})
         if sort is not None:
             params.update({"sort": sort})
         if direction is not None:
@@ -579,6 +609,8 @@ def list_models(
     def list_datasets(
         self,
         filter: Union[str, Iterable[str], None] = None,
+        author: Optional[str] = None,
+        search: Optional[str] = None,
         sort: Union[Literal["lastModified"], str, None] = None,
         direction: Optional[Literal[-1]] = None,
         limit: Optional[int] = None,
@@ -603,6 +635,30 @@ def list_datasets(
 
                     >>> # List only the datasets in russian for language modeling
                     >>> api.list_datasets(filter=("languages:ru", "task_ids:language-modeling"))
+            author (:obj:`str`, `optional`):
+                A string which identify the author of the returned models
+                Example usage:
+
+                    >>> from huggingface_hub import HfApi
+                    >>> api = HfApi()
+
+                    >>> # List all datasets from google
+                    >>> api.list_datasets(author="google")
+
+                    >>> # List only the text classification datasets from google
+                    >>> api.list_datasets(filter="text-classification", author="google")
+            search (:obj:`str`, `optional`):
+                A string that will be contained in the returned models
+                Example usage:
+
+                    >>> from huggingface_hub import HfApi
+                    >>> api = HfApi()
+
+                    >>> # List all datasets with "text" in their name
+                    >>> api.list_datasets(search="text")
+
+                    >>> #List all datasets with "text" in their name made by google
+                    >>> api.list_datasets(search="text", author="google")
             sort (:obj:`Literal["lastModified"]` or :obj:`str`, `optional`):
                 The key with which to sort the resulting datasets. Possible values are the properties of the `DatasetInfo`
                 class.
@@ -619,6 +675,10 @@ def list_datasets(
         params = {}
         if filter is not None:
             params.update({"filter": filter})
+        if author is not None:
+            params.update({"author": author})
+        if search is not None:
+            params.update({"search": search})
         if sort is not None:
             params.update({"sort": sort})
         if direction is not None:

tests/test_hf_api.py

@@ -472,6 +472,22 @@ def test_list_models(self):
         self.assertGreater(len(models), 100)
         self.assertIsInstance(models[0], ModelInfo)
 
+    @with_production_testing
+    def test_list_models_author(self):
+        _api = HfApi()
+        models = _api.list_models(author="google")
+        self.assertGreater(len(models), 10)
+        self.assertIsInstance(models[0], ModelInfo)
+        [self.assertTrue("google" in model.author for model in models)]
+
+    @with_production_testing
+    def test_list_models_search(self):
+        _api = HfApi()
+        models = _api.list_models(search="bert")
+        self.assertGreater(len(models), 10)
+        self.assertIsInstance(models[0], ModelInfo)
+        [self.assertTrue("bert" in model.modelId.lower()) for model in models]
+
     @with_production_testing
     def test_list_models_complex_query(self):
         # Let's list the 10 most recent models
@@ -549,6 +565,20 @@ def test_list_datasets_full(self):
         self.assertIsInstance(dataset, DatasetInfo)
         self.assertTrue(any(dataset.cardData for dataset in datasets))
 
+    @with_production_testing
+    def test_list_datasets_author(self):
+        _api = HfApi()
+        datasets = _api.list_datasets(author="huggingface")
+        self.assertGreater(len(datasets), 1)
+        self.assertIsInstance(datasets[0], DatasetInfo)
+
+    @with_production_testing
+    def test_list_datasets_search(self):
+        _api = HfApi()
+        datasets = _api.list_datasets(search="wikipedia")
+        self.assertGreater(len(datasets), 10)
+        self.assertIsInstance(datasets[0], DatasetInfo)
+
     @with_production_testing
     def test_dataset_info(self):
         _api = HfApi()