model config manager class

pyproject.toml

@@ -31,7 +31,9 @@ dependencies = [
     'tokenizers >= 0.13.2; python_version >= "3.11"',  # See https://github.com/citadel-ai/langcheck/pull/45
     'torch >= 2',
     'transformers >= 4.6',
-    "unidic-lite >= 1.0.1"  # For tokenizer of metrics.ja.toxicity()
+    "unidic-lite >= 1.0.1",  # For tokenizer of metrics.ja.toxicity()
+    "tabulate >= 0.9.0", # For model manager paint table
+    "omegaconf >= 2.3.0" # For model manager paint table
 ]
 requires-python = ">=3.8"
 
@@ -80,3 +82,9 @@ ignore = [
 markers = [
     "optional: marks tests as optional",
 ]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.package-data]
+langcheck = ["metrics/model_manager/config/*.yaml"]

src/langcheck/metrics/__init__.py

@@ -39,4 +39,5 @@
     'semantic_similarity',
     'sentiment',
     'toxicity',
+    'model_manager'
 ]

src/langcheck/metrics/model_manager/__init__.py

@@ -0,0 +1,11 @@
+from ._model_loader import (load_auto_model_for_seq2seq,
+                            load_auto_model_for_text_classification,
+                            load_sentence_transformers)
+from ._model_management import ModelManager
+
+manager = ModelManager()
+
+__all__ = [
+    "manager", "load_sentence_transformers", "load_auto_model_for_seq2seq",
+    "load_auto_model_for_text_classification"
+]

src/langcheck/metrics/model_manager/_model_loader.py

@@ -0,0 +1,87 @@
+from typing import Optional, Tuple
+
+from sentence_transformers import SentenceTransformer
+from transformers.models.auto.modeling_auto import (
+    AutoModelForSeq2SeqLM, AutoModelForSequenceClassification)
+from transformers.models.auto.tokenization_auto import AutoTokenizer
+
+
+def load_sentence_transformers(
+        model_name: str,
+        tokenizer_name: Optional[str] = None,
+        revision: Optional[str] = None) -> SentenceTransformer:
+    '''
+    Loads a SentenceTransformer model.
+
+    This function currently does not support specifying a tokenizer or a
+    revision. If these arguments are provided, a warning message will be
+    printed.
+
+    Args:
+        model_name: The name of the SentenceTransformer model to load.
+        tokenizer_name: The name of the tokenizer to use. Currently not
+            supported.
+        revision: The model revision to load. Currently not supported.
+
+    Returns:
+        model: The loaded SentenceTransformer model.
+    '''
+    if revision is not None:
+        print("Warning: Specifying a revision is not currently supported.")
+    if tokenizer_name is not None:
+        print("Warning: Customizing the tokenizer is not currently supported.")
+
+    model = SentenceTransformer(model_name)
+    return model
+
+
+def load_auto_model_for_text_classification(
+    model_name: str,
+    tokenizer_name: Optional[str] = None,
+    revision: Optional[str] = None
+) -> Tuple[AutoTokenizer, AutoModelForSequenceClassification]:
+    '''
+    Loads a sequence classification model and its tokenizer.
+
+    Args:
+        model_name: The name of the sequence-classification model to load.
+        tokenizer_name: The name of the tokenizer to load. If None, the
+            tokenizer associated with the model will be loaded.
+        revision: The model revision to load.
+
+    Returns:
+        tokenizer: The loaded tokenizer.
+        model: The loaded sequence classification model.
+    '''
+    if tokenizer_name is None:
+        tokenizer_name = model_name
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, revision=revision)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        model_name, revision=revision)
+    return tokenizer, model  # type: ignore
+
+
+def load_auto_model_for_seq2seq(
+    model_name: str,
+    tokenizer_name: Optional[str] = None,
+    revision: Optional[str] = None
+) -> Tuple[AutoTokenizer, AutoModelForSeq2SeqLM]:
+    '''
+    Loads a sequence-to-sequence model and its tokenizer.
+
+    Args:
+        model_name: The name of the sequence-classification model to load.
+        tokenizer_name: The name of the tokenizer to load. If None, the
+            tokenizer associated with the model will be loaded.
+        revision: The model revision to load.
+
+    Returns:
+        tokenizer: The loaded tokenizer.
+        model: The loaded sequence-to-sequence model.
+    '''
+    if tokenizer_name is None:
+        tokenizer_name = model_name
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, revision=revision)
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        model_name, revision=revision)  # NOQA: E501
+    return tokenizer, model  # type: ignore

src/langcheck/metrics/model_manager/_model_management.py

@@ -0,0 +1,257 @@
+import os
+from copy import deepcopy
+from functools import lru_cache
+from typing import Optional, Tuple, Union
+
+import pandas as pd
+import requests
+from omegaconf import OmegaConf
+from sentence_transformers import SentenceTransformer
+from tabulate import tabulate
+from transformers.models.auto.modeling_auto import (
+    AutoModelForSeq2SeqLM, AutoModelForSequenceClassification)
+from transformers.models.auto.tokenization_auto import AutoTokenizer
+
+from ._model_loader import (load_auto_model_for_seq2seq,
+                            load_auto_model_for_text_classification,
+                            load_sentence_transformers)
+
+LOADER_MAP = {
+    "load_sentence_transformers":
+        load_sentence_transformers,
+    "load_auto_model_for_text_classification":
+        load_auto_model_for_text_classification,
+    "load_auto_model_for_seq2seq":
+        load_auto_model_for_seq2seq
+}
+VALID_LOADER_FUNCTION = LOADER_MAP.keys()
+VALID_METRICS = [
+    'semantic_similarity', 'sentiment', 'toxicity', 'factual_consistency'
+]
+
+VALID_METRIC_ATTRIBUTE = [
+    'model_revision', 'model_revision', 'loader', 'tokenizer_name'
+]
+VALID_LANGUAGE = ['zh']
+
+
+def check_model_availability(model_name: str, revision: Optional[str]) -> bool:
+    # TODO: add local cached model availability check for offline environment
+    if revision is None:
+        url = f"https://huggingface.co/api/models/{model_name}"
+    else:
+        url = f"https://huggingface.co/api/models/{model_name}/revision/{revision}"  # NOQA:E501
+    response = requests.get(url, timeout=(1.0, 1.0))
+    return response.status_code == 200
+
+
+class ModelManager:
+    '''
+    A class to manage different models for multiple languages in LangCheck.
+    This class allows setting and retrieving different model names (like
+    sentiment_model, semantic_similarity_model, etc.) for each language.
+    It also supports loading model configurations from a file.
+    '''
+
+    def __init__(self):
+        '''
+        Initializes the ModelConfig with empty model dictionaries for each
+        language.
+        '''
+        self.config = OmegaConf.create()
+        cwd = os.path.dirname(__file__)
+        default_config_file_path = os.path.join(cwd, "config",
+                                                "metric_config.yaml")
+        self.__load_config(default_config_file_path)
+
+    def __load_config(self, path: str) -> None:
+        '''
+        Loads the model configuration from a file.
+
+        Args:
+            path: The path to the configuration file.
+        '''
+        conf = OmegaConf.load(path)
+
+        for lang, lang_conf in conf.items():
+            for metric_name, metric_conf in lang_conf.items():
+                # check model availbility, if key not in conf
+                # omega conf will return None in default
+                self.__set_model_for_metric(language=lang,
+                                            metric=metric_name,
+                                            **metric_conf)
+        print('Configuration Load Succeeded!')
+
+    @lru_cache
+    def fetch_model(
+        self, language: str, metric: str
+    ) -> Union[Tuple[AutoTokenizer, AutoModelForSequenceClassification], Tuple[
+            AutoTokenizer, AutoModelForSeq2SeqLM], SentenceTransformer]:
+        '''
+        Return the model (and if applicable, the tokenizer) used for the given
+        metric and language.
+
+        Args:
+            language: The language for which to get the model
+            metric_type: The metric name
+
+        Returns:
+            A (tokenizer, modle) tuple, or just the model depending on the
+            loader function.
+        '''
+        if language in self.config:
+            if metric in self.config[language]:
+                # Deep copy the confguration so that changes to `config` would
+                # not affect the original `self.config`.
+                config = deepcopy(self.config[language][metric])
+                # Get model loader function
+                loader_func = config.pop('loader_func')
+                loader = LOADER_MAP[loader_func]
+                # Call the loader function with the model_name, tokenizer_name
+                # (optional), and revision (optional) as arguments
+                return loader(**config)
+            else:
+                raise KeyError(f'Metric {metric} not supported yet.')
+        else:
+            raise KeyError(f'Language {language} not supported yet')
+
+    @staticmethod
+    def validate_config(config, language='all', metric='all') -> None:
+        '''
+        Validate configuration.
+
+        Args:
+            config: The configuration dictionary to validate.
+            language: The name of the language. Defaults to 'all'.
+            metric: The name of the metric. Defaults to 'all'.
+        '''
+        config = deepcopy(config)
+        for lang, lang_setting in config.items():
+            if language != 'all' and lang != language:
+                continue
+            for metric_name, model_setting in lang_setting.items():
+                if metric != 'all' and metric_name != metric:
+                    continue
+
+                # Check that the model name and loader function are set
+                if 'model_name' not in model_setting:
+                    raise KeyError(
+                        f'{lang} metrics {metric_name} need a model, but found None!'  # NOQA:E501
+                    )
+                if 'loader_func' not in model_setting:
+                    raise KeyError(
+                        f'Metrics {metric_name} need a loader, but found None!'  # NOQA:E501
+                    )
+                loader_func = model_setting.pop('loader_func', None)
+                if loader_func not in VALID_LOADER_FUNCTION:
+                    raise ValueError(
+                        f'loader type should in {VALID_LOADER_FUNCTION}')
+
+                # Check that the model and revision are available on the Hugging
+                # Face Hub
+                model_name = model_setting.pop('model_name')
+                revision = model_setting.pop('revision', None)
+                if not check_model_availability(model_name, revision):
+                    raise ValueError(
+                        f'Cannot find {model_name} with {revision} and Huggingface Hub'  # NOQA:E501
+                    )
+
+    def __set_model_for_metric(self, language: str, metric: str,
+                               model_name: str, loader_func: str,
+                               **kwargs) -> None:
+        '''
+        Set model for specified metric in specified language.
+
+        Args:
+            language: The name of the language
+            metric: The name of the evaluation metric
+            model_name: The name of the model
+            loader_func: The loader function of the model
+            tokenizer_name: (Optional) The name of the tokenizer
+            revision: (Optional) A version string of the model
+        '''
+        config_copy = deepcopy(self.config)
+        try:
+            if language not in VALID_LANGUAGE:
+                raise KeyError('Language {language} not supported yet')
+
+            if metric not in VALID_METRICS:
+                raise KeyError(
+                    f'Metric {metric} not supported for language {language} yet'
+                )
+
+            # Initialize the configuration for the language and metric if it
+            # doesn't exist
+            if self.config.get(language) is None:
+                self.config[language] = {}
+            if self.config.get(language).get(metric) is None:
+                self.config[language][metric] = {}
+
+            detail_config = self.config[language][metric]
+            # Set the loader function and model name
+            detail_config['loader_func'] = loader_func
+            detail_config['model_name'] = model_name
+
+            # If tokenizer_name is different from model_name
+            tokenizer_name = kwargs.pop('tokenizer_name', None)
+            if tokenizer_name:
+                detail_config['tokenizer_name'] = tokenizer_name
+            # If model's revision is pinned
+            revision = kwargs.pop('model_revision', None)
+            if revision:
+                detail_config['revision'] = revision
+
+            # Validate the change
+            if ModelManager.validate_config(self.config,
+                                            language=language,
+                                            metric=metric):
+                # Clear the LRU cache to make the config change reflected
+                # immediately
+                self.fetch_model.cache_clear()
+        except (ValueError, KeyError) as err:
+            # If an error occurred, restore the original configuration
+            self.config = config_copy
+            raise err
+
+    def list_current_model_in_use(self, language='all', metric='all') -> None:
+        '''
+        List the models currently in use.
+
+        Args:
+            language: The abbrevation name of language
+            metric: The evaluation metric name
+        '''
+        df = pd.DataFrame.from_records(
+            [(lang, metric_name, key, value)
+             for lang, lang_model_settings in self.config.items()
+             for metric_name, model_settings in lang_model_settings.items()
+             for key, value in model_settings.items()],
+            columns=['language', 'metric_name', 'attribute', 'value'])
+        # The code below would generate a dataframe:
+        # |index| language | metric_name | loader | model_name | revision |
+        # |.....|..........|.............|........|............|..........|
+        df_pivot = df.pivot_table(index=['language', 'metric_name'],
+                                  columns="attribute",
+                                  values="value",
+                                  aggfunc='first').reset_index().rename_axis(
+                                      None, axis=1)
+        df_pivot.columns = [
+            'language', 'metric_name', 'loader', 'model_name', 'revision'
+        ]
+
+        if language == 'all' and metric == 'all':
+            print(
+                tabulate(
+                    df_pivot,  # type: ignore
+                    headers=df_pivot.columns,  # type: ignore
+                    tablefmt="github"))
+        else:
+            if language != "all":
+                df_pivot = df_pivot.loc[df_pivot.language == language]
+            if metric != 'all':
+                df_pivot = df_pivot.loc[df_pivot.metric_name == metric]
+            print(
+                tabulate(
+                    df_pivot,  # type: ignore
+                    headers=df_pivot.columns,  # type: ignore
+                    tablefmt="github"))