def hf_token(hf_api: HfApi):
hf_token = hf_api.login(username=USER, password=PASS)
yield hf_token
try:
hf_api.logout(hf_token)
except requests.exceptions.HTTPError:
pass
class FlairModelHub:
"""
A class for interacting with the HF model hub API, and searching for Flair models by name or task
Can optionally include your HuggingFace login for authorized access (but is not required)
"""
def __init__(self, username=None, password=None):
self.api = HfApi()
if username and password:
self.token = self.api.login(username, password)
elif username or password:
print(
'Only a username or password was entered. You should include both to get authorized access'
)
self.models = self.api.list_models('flair') + FLAIR_MODELS
def _format_results(self,
results: list,
as_dict=False,
user_uploaded=False
) -> (List[HFModelResult], Dict[str, HFModelResult]):
"""
Takes raw HuggingFace API results and makes them easier to read and work with
"""
results = apply(FlairModelResult, results)
if not user_uploaded:
results = [
r for r in results
if 'flair/' in r.name or 'flairNLP/' in r.name
]
if as_dict:
dicts = apply(Self.to_dict(), results)
results = {m['model_name']: m for m in dicts}
return results
def search_model_by_name(
self,
name: str,
as_dict=False,
user_uploaded=False
) -> (List[HFModelResult], Dict[str, HFModelResult]):
"""
Searches HuggingFace Model API for all flair models containing `name` and returns a list of `HFModelResults`
Optionally can return all models as `dict` rather than a list
If `user_uploaded` is False, will only return models originating from Flair (such as flair/chunk-english-fast)
Usage:
```python
hub = FlairModelHubSearch()
hub.search_model_by_name('flair/chunk-english-fast')
```
"""
models = [m for m in self.models if name in m.modelId]
return self._format_results(models, as_dict, user_uploaded)
def search_model_by_task(
self,
task: str,
as_dict=False,
user_uploaded=False
) -> (List[HFModelResult], Dict[str, HFModelResult]):
"""
Searches HuggingFace Model API for all flair models for `task` and returns a list of `HFModelResults`
Optionally can return all models as `dict` rather than a list
If `user_uploaded` is False, will only return models originating from Flair (such as flair/chunk-english-fast)
Usage:
```python
hub = FlairModelHubSearch()
hub.search_model_by_task('ner')
# OR: #
hub.search_model_by_task(FLAIR_TASKS.NAMED_ENTITY_RECOGNITION)
```
"""
if (task not in _flair_tasks.values()) and (task != ''):
raise ValueError(f'''`{task}` is not a valid task.
Please choose a valid one available from Flair: (https://huggingface.co/flair)
Or with the `FLAIR_TASKS` object''')
models = [
m for m in self.models
if task in m.modelId or task == m.pipeline_tag
]
return self._format_results(models, as_dict, user_uploaded)
class HFModelHub:
"""
A class for interacting with the HF model hub API, and searching for models by name or task
Can optionally include your HuggingFace login for authorized access (but is not required)
"""
def __init__(self, username=None, password=None):
self.api = HfApi()
if username and password:
self.token = self.api.login(username, password)
elif username or password:
print(
'Only a username or password was entered. You should include both to get authorized access'
)
def _format_results(self,
results: list,
as_dict=False,
user_uploaded=False
) -> (List[HFModelResult], Dict[str, HFModelResult]):
"""
Takes raw HuggingFace API results and makes them easier to read and work with
"""
results = apply(HFModelResult, results)
if not user_uploaded:
results = [r for r in results if '/' not in r.name]
if as_dict:
dicts = apply(Self.to_dict(), results)
results = {m['model_name']: m for m in dicts}
return results
def search_model_by_task(
self,
task: str,
as_dict=False,
user_uploaded=False
) -> (List[HFModelResult], Dict[str, HFModelResult]):
"""
Searches HuggingFace Model API for all pretrained models relating to `task` and returns a list of HFModelResults
Optionally can return all models as a `dict` rather than a list
If `user_uploaded` is False, will only return models originating in HuggingFace (such as distilgpt2)
Usage:
```python
hub = HFModelHubSearch()
hub.search_model_by_task('summarization')
# OR #
hub.search_model_by_task(HF_TASKS.SUMMARIZATION)
```
"""
if task not in _hf_tasks.values():
raise ValueError(f'''`{task}` is not a valid task.
Please choose a valid one available from HuggingFace: (https://huggingface.co/transformers/task_summary.html)
Or with the `HF_TASKS` object''')
models = self.api.list_models(task)
return self._format_results(models, as_dict, user_uploaded)
def search_model_by_name(
self,
name: str,
as_dict=False,
user_uploaded=False
) -> (List[HFModelResult], Dict[str, HFModelResult]):
"""
Searches HuggingFace Model API for all pretrained models containing `name` and returns a list of HFModelResults
Optionally can return all models as `dict` rather than a list
If `user_uploaded` is False, will only return models originating from HuggingFace (such as distilgpt2)
Usage:
```python
hub = HFModelHubSearch()
hub.search_model_by_name('gpt2')
```
"""
if user_uploaded:
models = self.api.list_models()
models = self._format_results(models, as_dict, user_uploaded)
models = [m for m in models if name in m.name]
else:
models = self.api.list_models(name)
models = self._format_results(models, as_dict, user_uploaded)
return models
