def test_filter_models_by_language(self):
_api = HfApi()
f_fr = ModelFilter(language="fr")
res_fr = _api.list_models(filter=f_fr)
f_en = ModelFilter(language="en")
res_en = _api.list_models(filter=f_en)
assert len(res_fr) != len(res_en)
def test_filter_models_with_task(self):
_api = HfApi()
f = ModelFilter(task="fill-mask", model_name="albert-base-v2")
models = _api.list_models(filter=f)
self.assertTrue("fill-mask" == models[0].pipeline_tag)
self.assertTrue("albert-base-v2" in models[0].modelId)
f = ModelFilter(task="dummytask")
models = _api.list_models(filter=f)
self.assertGreater(1, len(models))
def test_filter_models_with_library(self):
_api = HfApi()
f = ModelFilter("microsoft",
model_name="wavlm-base-sd",
library="tensorflow")
models = _api.list_models(filter=f)
self.assertGreater(1, len(models))
f = ModelFilter("microsoft",
model_name="wavlm-base-sd",
library="pytorch")
models = _api.list_models(filter=f)
self.assertGreater(len(models), 0)
def test_filter_models_with_complex_query(self):
_api = HfApi()
args = ModelSearchArguments()
f = ModelFilter(
task=args.pipeline_tag.TextClassification,
library=[args.library.PyTorch, args.library.TensorFlow],
)
models = _api.list_models(filter=f)
self.assertGreater(len(models), 1)
self.assertTrue([
"text-classification" in model.pipeline_tag
or "text-classification" in model.tags for model in models
])
self.assertTrue([
"pytorch" in model.tags and "tf" in model.tags for model in models
])
def test_failing_filter_models_by_author_and_model_name(self):
# Test we can search by an author and a name, but the model is not found
_api = HfApi()
f = ModelFilter(author="muellerzr", model_name="testme")
models = _api.list_models(filter=f)
self.assertEqual(len(models), 0)
def test_filter_models_by_author_and_name(self):
# Test we can search by an author and a name, but the model is not found
_api = HfApi()
f = ModelFilter("facebook", model_name="bart-base")
models = _api.list_models(filter=f)
self.assertTrue("facebook/bart-base" in models[0].modelId)
def test_filter_models_by_author(self):
_api = HfApi()
f = ModelFilter(author="muellerzr")
models = _api.list_models(filter=f)
self.assertGreater(len(models), 0)
self.assertTrue("muellerzr" in models[0].modelId)