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,69 @@
+from typing import Optional, Tuple
+
+from sentence_transformers import SentenceTransformer
+from transformers.models.auto.modeling_auto import (  # NOQA:E501
+    AutoModelForSeq2SeqLM, AutoModelForSequenceClassification)
+from transformers.models.auto.tokenization_auto import AutoTokenizer
+from transformers.pipelines import pipeline
+
+
+def load_sentence_transformers(
+        model_name: str,
+        tokenizer_name: Optional[str] = None,
+        revision: Optional[str] = None) -> SentenceTransformer:  # NOQA:E501
+    '''
+    Return a sequence embeddiing model parsed by sentence-transformer library.
+
+    Args:
+        model_name: The name of a sentence-transformer model
+    '''
+    if revision is not None:
+        print("Version Pined not supported in Sentence-Transformers yet.")
+
+    if tokenizer_name is not None:
+        print("Tokenizer customize not supported in Sentence-Transformers yet.")
+
+    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]:
+    '''
+    Return a sequence classification model on huggingface hub.
+
+    Args:
+        model_name: The name of a sequence-classification model on Hugging Face
+        tokenizer_name: The name of a tokenizer on Hugging Face
+        revision: The shortened sha1 string of a 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)  # NOQA: E501
+    return tokenizer, model
+
+
+def load_auto_model_for_seq2seq(
+    model_name: str,
+    tokenizer_name: Optional[str] = None,
+    revision: Optional[str] = None
+) -> Tuple[AutoTokenizer, AutoModelForSequenceClassification]:
+    '''
+    Return a sequence to sequence model availble on huggingface hub.
+
+    Args:
+        model_name: The name of a sequence-classification model on Hugging Face
+        tokenizer_name: The name of a tokenizer on Hugging Face
+        revision: The shortened sha1 string of a 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

src/langcheck/metrics/model_manager/_model_management.py

@@ -0,0 +1,231 @@
+import os
+from copy import deepcopy
+from functools import lru_cache
+from pathlib import Path
+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 (  # NOQA:E501
+    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,  # NOQA:E501
+    "load_auto_model_for_seq2seq":
+        load_auto_model_for_seq2seq
+}
+VALID_LOADER_FUNCTION = LOADER_MAP.keys()  # NOQA:E501
+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]):
+    # 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):
+        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,   # type: ignore  # NOQA:E501
+                                            metric=metric_name,
+                                            **metric_conf)
+        print('Configuration Load Successed!')
+
+    @lru_cache
+    def fetch_model(
+        self, language: str, metric: str
+    ) -> Union[Tuple[AutoTokenizer, AutoModelForSequenceClassification], Tuple[
+            AutoTokenizer, AutoModelForSeq2SeqLM], SentenceTransformer]:
+        '''
+        Return the model used for the given metric and language.
+
+        Args:
+            language: The language for which to get the model
+            metric_type: The metric name
+        '''
+        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 name, model loader type
+                loader_func = config.pop('loader_func')
+                loader = LOADER_MAP[loader_func]
+                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'):
+        '''
+        Validate configuration.
+
+        Args:
+            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
+                # If model name not 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
+                    )
+                # Check if the model and revision is available on
+                # Hugging Face Hub
+                model_name = model_setting.pop('model_name')
+                revision = model_setting.pop('revision', None)
+                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}')
+                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):
+        '''
+        Set model for specified metric in specified language.
+
+        Args:
+            language: The name of the language
+            metric: The name of the evaluation metrics
+            model_name: The name of the model
+            loader: The loader 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('Language {language} not supported {metric} yet')
+
+            # initialize configuration structure if it is empty.
+            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 metric attribute
+            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:
+            # Trace back the configuration
+            self.config = config_copy
+            raise err
+
+    def list_current_model_in_use(self, language='all', metric='all'):
+        '''
+        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, headers=df_pivot.columns, tablefmt="github"))  # type: ignore  # NOQA:E501
+        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, headers=df_pivot.columns,  # type: ignore  # NOQA:E501
+                         tablefmt="github"))

src/langcheck/metrics/model_manager/config/metric_config.yaml

@@ -0,0 +1,19 @@
+zh:
+  semantic_similarity:
+    model_name: BAAI/bge-base-zh-v1.5
+    model_revision: f03589c
+    loader_func: load_sentence_transformers
+
+  sentiment:
+    model_name: IDEA-CCNL/Erlangshen-Roberta-110M-Sentiment
+    loader_func: load_auto_model_for_text_classification
+
+  toxicity:
+    model_name: alibaba-pai/pai-bert-base-zh-llm-risk-detection
+    model_revision: 0a61c79744cb0173216f015ffecc1ea81c4e0229
+    loader_func: load_auto_model_for_text_classification
+
+  factual_consistency:
+    model_name: Helsinki-NLP/opus-mt-zh-en
+    model_revision: cf109095479db38d6df799875e34039d4938aaa6
+    loader_func: load_auto_model_for_seq2seq

src/langcheck/metrics/zh/reference_based_text_quality.py

@@ -90,14 +90,13 @@ def semantic_similarity(
                                               openai_args)
         metric_value.language = 'zh'
         return metric_value
+    # lazy import
+    from langcheck.metrics.model_manager import manager
+    model = manager.fetch_model(language='zh',
+                                metric="semantic_similarity")
 
-    # According to the C-MTEB Benchmark
-    # (https://github.com/FlagOpen/FlagEmbedding/tree/master/C_MTEB)
-    # the 3 models of different sizes provided BAAI are the best on the
-    # embedding task
-    # Ref: https://huggingface.co/BAAI/bge-base-zh-v1.5
-    # Using this model, it is hard to find two sentence where cos_sim < 0.25.
-    model = SentenceTransformer('BAAI/bge-base-zh-v1.5')
+    # For type checking
+    assert isinstance(model, SentenceTransformer)
     generated_embeddings = model.encode(generated_outputs)
     reference_embeddings = model.encode(reference_outputs)
     cosine_scores = util.pairwise_cos_sim(