Source code for malaya.emotion

from malaya.supervised import classification
from malaya.path import PATH_EMOTION, S3_PATH_EMOTION
from herpetologist import check_type
from malaya.function import describe_availability
import logging

logger = logging.getLogger(__name__)

label = ['anger', 'fear', 'happy', 'love', 'sadness', 'surprise']
label_goemotions = [
    'admiration',
    'amusement',
    'anger',
    'annoyance',
    'approval',
    'caring',
    'confusion',
    'curiosity',
    'desire',
    'disappointment',
    'disapproval',
    'disgust',
    'embarrassment',
    'excitement',
    'fear',
    'gratitude',
    'grief',
    'joy',
    'love',
    'nervousness',
    'optimism',
    'pride',
    'realization',
    'relief',
    'remorse',
    'sadness',
    'surprise',
    'neutral'
]

_transformer_availability = {
    'bert': {
        'Size (MB)': 425.6,
        'Quantized Size (MB)': 111,
        'macro precision': 0.99786,
        'macro recall': 0.99773,
        'macro f1-score': 0.99779,
    },
    'tiny-bert': {
        'Size (MB)': 57.4,
        'Quantized Size (MB)': 15.4,
        'macro precision': 0.99692,
        'macro recall': 0.99696,
        'macro f1-score': 0.99694,
    },
    'albert': {
        'Size (MB)': 48.6,
        'Quantized Size (MB)': 12.8,
        'macro precision': 0.99740,
        'macro recall': 0.99773,
        'macro f1-score': 0.99757,
    },
    'tiny-albert': {
        'Size (MB)': 22.4,
        'Quantized Size (MB)': 5.98,
        'macro precision': 0.99325,
        'macro recall': 0.99378,
        'macro f1-score': 0.99351,
    },
    'xlnet': {
        'Size (MB)': 446.5,
        'Quantized Size (MB)': 118,
        'macro precision': 0.99773,
        'macro recall': 0.99775,
        'macro f1-score': 0.99774,
    },
    'alxlnet': {
        'Size (MB)': 46.8,
        'Quantized Size (MB)': 13.3,
        'macro precision': 0.99663,
        'macro recall': 0.99697,
        'macro f1-score': 0.99680,
    },
    'fastformer': {
        'Size (MB)': 446,
        'Quantized Size (MB)': 113,
        'macro precision': 0.99197,
        'macro recall': 0.99194,
        'macro f1-score': 0.99195,
    },
    'tiny-fastformer': {
        'Size (MB)': 77.2,
        'Quantized Size (MB)': 19.6,
        'macro precision': 0.98926,
        'macro recall': 0.98783,
        'macro f1-score': 0.98853,
    }
}


def available_transformer():
    """
    List available transformer emotion analysis models.
    """
    logger.info('trained on 80% dataset, tested on another 20% test set, dataset at https://github.com/huseinzol05/malay-dataset/tree/master/corpus/emotion')

    return describe_availability(_transformer_availability)


def multinomial(**kwargs):
    """
    Load multinomial emotion model.

    Returns
    -------
    result: malaya.model.ml.MulticlassBayes class
    """
    return classification.multinomial(
        path=PATH_EMOTION,
        s3_path=S3_PATH_EMOTION,
        module='emotion',
        label=label,
        **kwargs
    )


@check_type
def transformer(model: str = 'xlnet', quantized: bool = False, **kwargs):
    """
    Load Transformer emotion model.

    Parameters
    ----------
    model: str, optional (default='bert')
        Check available models at `malaya.emotion.available_transformer()`.
    quantized: bool, optional (default=False)
        if True, will load 8-bit quantized model.
        Quantized model not necessary faster, totally depends on the machine.

    Returns
    -------
    result: model
        List of model classes:

        * if `bert` in model, will return `malaya.model.bert.MulticlassBERT`.
        * if `xlnet` in model, will return `malaya.model.xlnet.MulticlassXLNET`.
        * if `fastformer` in model, will return `malaya.model.fastformer.MulticlassFastFormer`.
    """

    model = model.lower()
    if model not in _transformer_availability:
        raise ValueError(
            'model not supported, please check supported models from `malaya.emotion.available_transformer()`.'
        )
    return classification.transformer(
        module='emotion',
        label=label,
        model=model,
        quantized=quantized,
        **kwargs
    )