# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/02_model_hub.ipynb (unless otherwise specified).

__all__ = ['HF_TASKS', 'FLAIR_TASKS', 'HFModelResult', 'HFModelHub', 'FLAIR_MODELS', 'FlairModelResult',
           'FlairModelHub']

# Cell
from fastcore.basics import Self, merge
from fastcore.utils import dict2obj, obj2dict, mk_class
from fastai.torch_core import apply
from huggingface_hub.hf_api import ModelInfo, HfApi

from typing import List, Dict

# Internal Cell
_hf_tasks = {
    'FILL_MASK':'fill-mask',
    'QUESTION_ANSWERING':'question-answering',
    'SUMMARIZATION':'summarization',
    'TABLE_QUESTION_ANSWERING':'table-question-answering',
    'TEXT_CLASSIFICATION':'text-classification',
    'TEXT_GENERATION':'text-generation',
    'TEXT2TEXT_GENERATION':'text2text-generation',
    'TOKEN_CLASSIFICATION':'token-classification',
    'TRANSLATION':'translation',
    'ZERO_SHOT_CLASSIFICATION':'zero-shot-classification',
    'CONVERSATIONAL':'conversational',
    'TEXT_TO_SPEECH':'text-to-speech',
    'AUTOMATIC_SPEECH_RECOGNITION':'automatic-speech-recognition',
    'AUDIO_SOURCE_SEPERATION':'audio-source-seperation',
    'VOICE_ACTIVITY_DETECTION':'voice-activity-detection'}

# Cell
mk_class('HF_TASKS', **_hf_tasks,
        doc="A list of all HuggingFace tasks for valid API lookup as attribtues to get tab-completion and typo-proofing")

# Cell
#nbdev_comment _all_ = ['HF_TASKS']

# Internal Cell
_flair_tasks = {
    'NAMED_ENTITY_RECOGNITION':'ner',
    'PHRASE_CHUNKING':'chunk',
    'VERB_DISAMBIGUATION':'frame',
    'PART_OF_SPEECH_TAGGING':'pos',
    'UNIVERSAL_PART_OF_SPEECH_TAGGING':'upos',
    'EMBEDDINGS':'embeddings',

}

# Cell
mk_class('FLAIR_TASKS', **_flair_tasks,
        doc="A list of all Flair tasks for valid API lookup as attributes to get tab-completion and typo-proofing")

# Cell
#nbdev_comment _all_ = ['FLAIR_TASKS']

# Cell
class HFModelResult:
    """
    A very basic class for storing a HuggingFace model returned through an API request

    They have 4 properties:
      - `name`: The `modelId` from the `modelInfo`. This also includes the model author's name, such as "IlyaGusev/mbart_ru_sum_gazeta"
      - `tags`: Any tags that were included in `HuggingFace` in relation to the model.
      - `tasks`: These are the tasks dictated for the model.
    """
    def __init__(
        self,
        model_info: ModelInfo # `ModelInfo` object from HuggingFace model hub
    ):
        self.info = model_info

    def __repr__(self): return f'Model Name: {self.name}, Tasks: [' + ', '.join(self.tasks) + ']'

    @property
    def name(self): return self.info.modelId

    @property
    def tags(self): return self.info.tags

    @property
    def tasks(self):
        if self.info.pipeline_tag:
            all_tasks = [self.info.pipeline_tag]
            for tag in self.tags:
                if (tag in _hf_tasks.values()) and (tag not in all_tasks):
                    all_tasks += [tag]
        else: all_tasks = []
        all_tasks.sort()
        return all_tasks

    def to_dict(
        self
    ) -> dict: # Dictionary with keys `model_name`, `tags`, `tasks`, `model_info`
        "Returns `HFModelResult` as a dictionary"
        return {'model_name':self.name, 'tags':self.tags, 'tasks':self.tasks, 'model_info':self.info}

# Cell
class HFModelHub:
    "A class for interacting with the HF model hub API, and searching for models by name or task"

    def __init__(
        self,
        username:str=None, # Your HuggingFace username
        password:str=None # Your HuggingFace password
    ):
        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, # A list of HuggingFace API results
        as_dict:bool=False, # Whether to return as a dictionary or list
        user_uploaded:bool=False # Whether to filter out user-uploaded results
    ) -> (List[HFModelResult], Dict[str, HFModelResult]): # A list of `HFModelResult`s
        "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, # A valid task to search in the HuggingFace hub for
        as_dict:bool=False, # Whether to return as a dictionary or list
        user_uploaded:bool=False # Whether to filter out user-uploaded results
    ) -> (List[HFModelResult], Dict[str, HFModelResult]): # A list of `HFModelResult`s
        "Searches HuggingFace Model API for all pretrained models relating to `task`"
        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, # A valid model name
        as_dict:bool=False, # Whether to return as a dictionary or list
        user_uploaded:bool=False # Whether to filter out user-uploaded results
    ) -> (List[HFModelResult], Dict[str, HFModelResult]): # A list of `HFModelResult`s
        "Searches HuggingFace Model API for all pretrained models containing `name`"
        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

# Internal Cell
# Flair models originating from:
# https://github.com/flairNLP/flair/blob/master/flair/models/text_classification_model.py#L483
# and: https://github.com/flairNLP/flair/blob/master/flair/models/sequence_tagger_model.py#L1053
# and: https://github.com/flairNLP/flair/blob/master/flair/embeddings/token.py#L406
_flair_models = {
    'de-offensive-language' : ['text-classification'],
    'sentiment' : ['text-classification'],
    'en-sentiment' : ['text-classification'],
    'sentiment-fast' : ['text-classification'],
    'communicative-functions' : ['text-classification'],
    'tars-base' : ['text-classification'],
    # English Named Entity Recognition Models (NER)
    'ner' : ['token-classification'],
    'ner-pooled' : ['token-classification'],
    'ner-fast' : ['token-classification'],
    'ner-ontonotes' : ['token-classification'],
    'ner-ontonotes-fast' : ['token-classification'],
    # Multilingual NER models
    'ner-multi' : ['token-classification'],
    'multi-ner' : ['token-classification'],
    'ner-multi-fast' : ['token-classification'],
    # English POS models
    'upos' : ['token-classification'],
    'upos-fast' : ['token-classification'],
    'pos' : ['token-classification'],
    'pos-fast' : ['token-classification'],
    # Multilingual POS models
    'pos-multi' : ['token-classification'],
    'multi-pos' : ['token-classification'],
    'pos-multi-fast' : ['token-classification'],
    'multi-pos-fast' : ['token-classification'],
    # English SRL models
    'frame' : ['token-classification'],
    'frame-fast' : ['token-classification'],
    # English chunking models
    'chunk' : ['token-classification'],
    'chunk-fast' : ['token-classification'],
    # Danish models
    'da-pos' : ['token-classification'],
    'da-ner' : ['token-classification'],
    # German models
    'de-pos' : ['token-classification'],
    'de-pos-tweets' : ['token-classification'],
    'de-ner' : ['token-classification'],
    'de-ner-germeval' : ['token-classification'],
    'de-ler' : ['token-classification'],
    'de-ner-legal' : ['token-classification'],
    # French models
    'fr-ner' : ['token-classification'],
    # Dutch models
    'nl-ner' : ['token-classification'],
    'nl-ner-rnn' : ['token-classification'],
    # Malayalam models
    'ml-pos' : ['token-classification'],
    'ml-upos' : ['token-classification'],
    # Portuguese models
    'pt-pos-clinical' : ['token-classification'],
    # Keyphrase models
    'keyphrase' : ['token-classification'],
    'negation-speculation' : ['token-classification'],
    # Biomedical
    'hunflair-paper-cellline' : ['token-classification'],
    'hunflair-paper-chemical' : ['token-classification'],
    'hunflair-paper-disease' : ['token-classification'],
    'hunflair-paper-gene' : ['token-classification'],
    'hunflair-paper-species' : ['token-classification'],
    'hunflair-cellline' : ['token-classification'],
    'hunflair-chemical' : ['token-classification'],
    'hunflair-disease' : ['token-classification'],
    'hunflair-gene' : ['token-classification'],
    'hunflair-species' : ['token-classification'],
    # Embeddings
    # multilingual models
    "multi-forward":['embeddings'],
    "multi-backward":['embeddings'],
    "multi-v0-forward":['embeddings'],
    "multi-v0-backward":['embeddings'],
    "multi-forward-fast":['embeddings'],
    "multi-backward-fast":['embeddings'],
    # English models
    "en-forward":['embeddings'],
    "en-backward":['embeddings'],
    "en-forward-fast":['embeddings'],
    "en-backward-fast":['embeddings'],
    "news-forward":['embeddings'],
    "news-backward":['embeddings'],
    "news-forward-fast":['embeddings'],
    "news-backward-fast":['embeddings'],
    "mix-forward":['embeddings'],
    "mix-backward":['embeddings'],
    # Arabic
    "ar-forward":['embeddings'],
    "ar-backward":['embeddings'],
    # Bulgarian
    "bg-forward-fast":['embeddings'],
    "bg-backward-fast":['embeddings'],
    "bg-forward":['embeddings'],
    "bg-backward":['embeddings'],
    # Czech
    "cs-forward":['embeddings'],
    "cs-backward":['embeddings'],
    "cs-v0-forward":['embeddings'],
    "cs-v0-backward":['embeddings'],
    # Danish
    "da-forward":['embeddings'],
    "da-backward":['embeddings'],
    # German
    "de-forward":['embeddings'],
    "de-backward":['embeddings'],
    "de-historic-ha-forward":['embeddings'],
    "de-historic-ha-backward":['embeddings'],
    "de-historic-wz-forward":['embeddings'],
    "de-historic-wz-backward":['embeddings'],
    "de-historic-rw-forward":['embeddings'],
    "de-historic-rw-backward":['embeddings'],
    # Spanish
    "es-forward":['embeddings'],
    "es-backward":['embeddings'],
    "es-forward-fast":['embeddings'],
    "es-backward-fast":['embeddings'],
    # Basque
    "eu-forward":['embeddings'],
    "eu-backward":['embeddings'],
    "eu-v1-forward":['embeddings'],
    "eu-v1-backward":['embeddings'],
    "eu-v0-forward":['embeddings'],
    "eu-v0-backward":['embeddings'],
    # Persian
    "fa-forward":['embeddings'],
    "fa-backward":['embeddings'],
    # Finnish
    "fi-forward":['embeddings'],
    "fi-backward":['embeddings'],
    # French
    "fr-forward":['embeddings'],
    "fr-backward":['embeddings'],
    # Hebrew
    "he-forward":['embeddings'],
    "he-backward":['embeddings'],
    # Hindi
    "hi-forward":['embeddings'],
    "hi-backward":['embeddings'],
    # Croatian
    "hr-forward":['embeddings'],
    "hr-backward":['embeddings'],
    # Indonesian
    "id-forward":['embeddings'],
    "id-backward":['embeddings'],
    # Italian
    "it-forward":['embeddings'],
    "it-backward":['embeddings'],
    # Japanese
    "ja-forward":['embeddings'],
    "ja-backward":['embeddings'],
    # Malayalam
    "ml-forward":['embeddings'],
    "ml-backward":['embeddings'],
    # Dutch
    "nl-forward":['embeddings'],
    "nl-backward":['embeddings'],
    "nl-v0-forward":['embeddings'],
    "nl-v0-backward":['embeddings'],
    # Norwegian
    "no-forward":['embeddings'],
    "no-backward":['embeddings'],
    # Polish
    "pl-forward":['embeddings'],
    "pl-backward":['embeddings'],
    "pl-opus-forward":['embeddings'],
    "pl-opus-backward":['embeddings'],
    # Portuguese
    "pt-forward":['embeddings'],
    "pt-backward":['embeddings'],
    # Pubmed
    "pubmed-forward":['embeddings'],
    "pubmed-backward":['embeddings'],
    "pubmed-2015-forward":['embeddings'],
    "pubmed-2015-backward":['embeddings'],
    # Slovenian
    "sl-forward":['embeddings'],
    "sl-backward":['embeddings'],
    "sl-v0-forward":['embeddings'],
    "sl-v0-backward":['embeddings'],
    # Swedish
    "sv-forward":['embeddings'],
    "sv-backward":['embeddings'],
    "sv-v0-forward":['embeddings'],
    "sv-v0-backward":['embeddings'],
    # Tamil
    "ta-forward":['embeddings'],
    "ta-backward":['embeddings'],
    # CLEF HIPE Shared task
    "de-impresso-hipe-v1-forward":['embeddings'],
    "de-impresso-hipe-v1-backward":['embeddings'],
    "en-impresso-hipe-v1-forward":['embeddings'],
    "en-impresso-hipe-v1-backward":['embeddings'],
    "fr-impresso-hipe-v1-forward":['embeddings'],
    "fr-impresso-hipe-v1-backward":['embeddings']
}

# Cell
FLAIR_MODELS = [ModelInfo(f'flairNLP/{key}', pipeline_tag=val[0]) for key,val in _flair_models.items()]

# Cell
class FlairModelResult(HFModelResult):
    """
    A version of `HFModelResult` for Flair specifically.

    Includes which backend the model was found (such as on HuggingFace or Flair's private model list)
    """

    def __init__(
        self,
        model_info: ModelInfo # ModelInfo object from HuggingFace model hub
    ):
        if 'flairNLP' in model_info.modelId:
            self.from_hf = False
        else:
            self.from_hf = True
        super().__init__(model_info)

    def __repr__(self): return f'Model Name: {self.name.replace("flairNLP", "flair")}, Tasks: [' + ', '.join(self.tasks) + ']' + f', Source: {self.source}'

    @property
    def source(self):
        if self.from_hf: return "HuggingFace Model Hub"
        else: return "Flair's Private Model Hub"

# Cell
class FlairModelHub:
    "A class for interacting with the HF model hub API, and searching for Flair models by name or task"

    def __init__(
        self,
        username:str=None, # HuggingFace username
        password:str=None # HuggingFace password
    ):
        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, # A list of HuggingFace API results
        as_dict:bool=False, # Whether to return as a dictionary or list
        user_uploaded:bool=False # Whether to filter out user-uploaded results
    ) -> (List[FlairModelResult], Dict[str, FlairModelResult]): # A list of `FlairModelResult`s
        "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, # A valid model name
        as_dict:bool=False, # Whether to return as a dictionary or list
        user_uploaded:bool=False # Whether to filter out user-uploaded results
    ) -> (List[FlairModelResult], Dict[str, FlairModelResult]): # A list of `FlairModelResult`s
        "Searches HuggingFace Model API for all flair models containing `name`"
        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, # A valid task to search the HuggingFace hub for
        as_dict=False, # Whether to return as a dictionary or list
        user_uploaded=False # Whether to filter out user-uploaded results
    ) -> (List[FlairModelResult], Dict[str, FlairModelResult]): # A list of `FlairModelResult`s
        "Searches HuggingFace Model API for all flair models for `task`"
        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)