model config manager class

src/langcheck/metrics/__init__.py

@@ -1,4 +1,5 @@
 from langcheck.metrics import en, ja, zh
+from langcheck.metrics._model_management import ModelManager
 from langcheck.metrics.en.reference_based_text_quality import (
     rouge1, rouge2, rougeL, semantic_similarity)
 from langcheck.metrics.en.reference_free_text_quality import (
@@ -13,6 +14,8 @@
                                               is_json_array, is_json_object,
                                               matches_regex, validation_fn)
 
+_model_manager = ModelManager()
+
 __all__ = [
     'en',
     'ja',

src/langcheck/metrics/_model_loader.py

@@ -0,0 +1,46 @@
+from typing import Optional, Tuple
+
+from sentence_transformers import SentenceTransformer
+from transformers.models.auto.modeling_auto import \
+    AutoModelForSequenceClassification
+from transformers.models.auto.tokenization_auto import AutoTokenizer
+from transformers.pipelines import pipeline
+
+
+def load_sentence_transformers(model_name: str) -> SentenceTransformer:
+    """
+    return a sentence-transformer model.
+
+    Args:
+        model_name: The model name of a sentence-transformers model
+    """
+    return SentenceTransformer(model_name)
+
+
+def load_auto_model_for_text_classification(model_name: str,
+                                            tokenizer_name: Optional[str],
+                                            revision: Optional[str])\
+                            -> Tuple[AutoTokenizer,
+                                     AutoModelForSequenceClassification]:
+    """
+    return a Huggingface text-classification pipeline.
+
+    Args:
+        model_name: The name of a sequenceclassification model on huggingface hub.  # NOQA:E501
+        tokenizer_name: the name of a tokenizer on huggingface hub.
+        revisoin: the shorted sha1 string of a model
+    """
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, revision=revision)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, revision=revision)  # NOQA: E501
+    return tokenizer, model
+
+
+def load_pipeline_for_text_classification(model_name: str, **kwargs):
+    """
+    return a Huggingface text-classification pipeline.
+
+    Args:
+        model_name: A huggingface model model for text classification.
+    """
+    top_k = kwargs.pop('top_k', None)
+    return pipeline('text-classification', model=model_name, top_k=top_k)

src/langcheck/metrics/_model_management.py

@@ -0,0 +1,226 @@
+import os
+from copy import deepcopy
+from functools import lru_cache
+from pathlib import Path
+from pprint import pprint
+from typing import Optional, Tuple, Union
+
+import pandas as pd
+import requests
+from configobj import ConfigObj
+from sentence_transformers import SentenceTransformer
+from transformers.models.auto.modeling_auto import \
+    AutoModelForSequenceClassification
+from transformers.models.auto.tokenization_auto import AutoTokenizer
+
+from ._model_loader import (load_auto_model_for_text_classification,
+                            load_sentence_transformers)
+
+# TODO: Use a ENUM class to parse these
+VALID_METRIC_NAME = [
+    'factual_consistency', 'toxicity', 'sentiment', 'semantic_similarity'
+]
+VALID_LANGUAGE = ['zh']
+VALID_LOADER = ['huggingface', 'sentence-transformers']
+
+
+class ModelManager:
+    """
+    A class to manage different models for multiple languages in the
+    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 = None
+        self.__init__config()
+        self.validate_config()
+
+    def __init__config(self):
+        cwd = os.path.dirname(__file__)
+        self.config = ConfigObj(
+            os.path.join(Path(cwd), 'config', 'metric_config.ini'))  # NOQA:E501
+
+    @lru_cache
+    def fetch_model(self, language: str, metric: str)\
+        -> Union[Tuple[AutoTokenizer, AutoModelForSequenceClassification],
+                 SentenceTransformer]:
+        """
+        return the model used in current metric for a given language.
+
+        Args:
+            language: The language for which to get the model.
+            metric_type: The metric name.
+        """
+        if language in self.config:  # type: ignore
+            if metric in self.config[language]:  # type: ignore
+                # deep copy the confguration
+                # any action on config would not distrub self.config
+                config = deepcopy(
+                    self.config[language][metric]  # type: ignore[reportGeneralTypeIssues]  # NOQA:E501
+                )
+                # get model name, model loader type
+                model_name, loader_type = config['model_name'], config[
+                    'loader']  # type: ignore[reportGeneralTypeIssues]  # NOQA:E501
+                # check if model version fixed
+                revision = config.pop("revision", None)
+                if loader_type == 'sentence-transformers':
+                    if revision is not None:
+                        print(
+                            'Info: Sentence-Transformers do not support model version fixed yet'  # NOQA: E501
+                        )
+                    model = load_sentence_transformers(model_name=model_name)
+                    return model
+                elif loader_type == 'huggingface':
+                    tokenizer_name = config.pop('tokenizer_name', None)
+                    tokenizer, model = load_auto_model_for_text_classification(
+                        model_name=model_name,  # NOQA:E501
+                        tokenizer_name=tokenizer_name,  # NOQA:E501
+                        revision=revision  # NOQA:E501
+                    )
+                    return tokenizer, model
+                else:
+                    raise KeyError(f'Loader {loader_type} not supported yet.')
+            else:
+                raise KeyError(f'Metric {metric} not supported yet.')
+        else:
+            raise KeyError(f'language {language} not supported yet')
+
+    def list_current_model_in_use(self, language='all', metric='all'):
+        """ list model 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()  # type: ignore  # NOQA:E501
+                for metric_name, model_settings in
+                lang_model_settings.items()  # type: ignore  # NOQA:E501
+                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().drop(
+                columns=["attribute"]).reset_index()  # NOQA:E501
+        df_pivot.columns = [
+            'language', 'metric_name', 'loader', 'model_name', 'revision'
+        ]  # NOQA:E501
+
+        if language == 'all' and metric == 'all':
+            pprint(df_pivot)
+        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]
+            pprint(df_pivot)
+
+    def validate_config(self, language='all', metric='all'):
+        """validate configuration.
+
+        Args:
+            language (str, optional):the name of the language. Defaults to 'all'.  # NOQA:E501
+            metric (str, optional): the name of evaluation metric. Defaults to 'all'.  # NOQA:E501
+        """
+
+        def check_model_availability(model_name, revision):
+            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)
+            return response.status_code == 200
+
+        config = deepcopy(self.config)
+        for lang, lang_setting in config.items():  # type: ignore  # NOQA:E501
+            if language == 'all' or lang == language:
+                for metric_name, model_setting in lang_setting.items(  # type: ignore  # NOQA:E501
+                ):
+                    if metric == 'all' or metric_name == metric:
+                        # 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' 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 huggingface Hub  # NOQA:E501
+                        loader_type = model_setting.pop('loader')
+                        if loader_type == 'huggingface':
+                            model_name = model_setting.pop('model_name')
+                            revision = model_setting.pop('revision', None)
+                            if not check_model_availability(
+                                    model_name, revision):  # NOQA:E501
+                                raise ValueError(
+                                    f"""Cannot find {model_name} with  # NOQA:E501
+                                                {revision} and Huggingface Hub"""
+                                )
+                        elif loader_type not in VALID_LOADER:
+                            raise ValueError(
+                                f'loader type should in {VALID_LOADER}'
+                            )  # NOQA: E501
+                        # may also need other validate method for other loader
+                        # not found yet
+        print('Configuration Validation Passed')
+
+    def set_model_for_metric(self, language: str, metric: str, model_name: str,
+                             loader: Optional[str], **kwargs):
+        """set model for specified metric in specified language
+
+        Args:
+            language (str): the name of the lanuage,
+            metric (str): the name of the evaluation metrics,
+            loader(str): the loader of the model, optional,
+            model_name(str): the name of the model,
+            tokenizer_name(str): optional, the name of the tokenizer,
+            revision(str): a version string of the model.
+        """
+        config_copy = deepcopy(self.config)
+        try:
+            if language not in VALID_LANGUAGE:
+                raise ValueError('Language {language} not supported yet')
+
+            if metric not in self.config[language]:  # type: ignore  # NOQA:E501
+                raise ValueError(
+                    'Language {language} not supported {metric} yet'
+                )  # NOQA:E501
+
+            config = self.config[language][metric]  # type: ignore  # NOQA:E501
+            config['loader'] = loader
+            config['model_name'] = model_name
+            # if tokenizer_name is different with model
+            tokenizer_name = kwargs.pop('tokenizer_name', None)
+            if tokenizer_name:
+                config['tokenizer_name'] = tokenizer_name
+            # if model's revision is pinned
+            revision = kwargs.pop('revision', None)
+            if revision:
+                config['revision'] = revision
+            # validate the change
+            if self.validate_config(language=language, metric=metric):
+                # clear the LRU cache to make the config change
+                # reflected imediately
+                self.fetch_model.cache_clear()
+        except (ValueError, KeyError) as err:
+            # trace back the configuration
+            self.config = config_copy
+            raise err

src/langcheck/metrics/config/metric_config.ini

@@ -0,0 +1,22 @@
+[zh]
+[[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_name = BAAI/bge-base-zh-v1.5
+    revision = f03589c
+    loader = sentence-transformers
+[[sentiment]]
+    model_name = IDEA-CCNL/Erlangshen-Roberta-110M-Sentiment
+    loader = huggingface
+[[toxicity]]
+    model_name = alibaba-pai/pai-bert-base-zh-llm-risk-detection
+    loader = huggingface
+    revision = 0a61c79744cb0173216f015ffecc1ea81c4e0229
+[[factual_consistency]]
+    model_name = Helsinki-NLP/opus-mt-zh-en
+    loader = huggingface
+    revision = cf109095479db38d6df799875e34039d4938aaa6

src/langcheck/metrics/zh/reference_based_text_quality.py

@@ -90,16 +90,13 @@ def semantic_similarity(
                                               openai_args)
         metric_value.language = 'zh'
         return metric_value
+    # lazy import
+    from langcheck.metrics import _model_manager
+    model = _model_manager.fetch_model(language='zh',
+                                       metric_type="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')
-    generated_embeddings = model.encode(generated_outputs)
-    reference_embeddings = model.encode(reference_outputs)
+    generated_embeddings = model.encode(generated_outputs)  # type: ignore[reportGeneralTypeIssues]  # NOQA: E501
+    reference_embeddings = model.encode(reference_outputs)  # type: ignore[reportGeneralTypeIssues]  # NOQA: E501
     cosine_scores = util.pairwise_cos_sim(
         generated_embeddings,  # type: ignore[reportGeneralTypeIssues]
         reference_embeddings  # type: ignore[reportGeneralTypeIssues]

src/langcheck/metrics/zh/reference_free_text_quality.py

@@ -92,7 +92,11 @@ def sentiment(
     _sentiment_pipeline = pipeline(
         'sentiment-analysis', model=_sentiment_model_path
     )  # type: ignore[reportGeneralTypeIssues]  # NOQA: E501
-    # {0:"Negative", 1:'Positive'}
+    # # {0:"Negative", 1:'Positive'}
+    from langcheck.metrics import _model_manager
+    tokenizer, model = _model_manager.fetch_model(lanaguage='zh', metric='sentiment')   # NOQA: E501
+    _sentiment_pipeline = pipeline(
+        'sentiment-analysis', model=model, tokenizer=tokenizer)  # type: ignore[reportGeneralTypeIssues]  # NOQA: E501
     _model_id2label = _sentiment_pipeline.model.config.id2label
     _predict_result = _sentiment_pipeline(
         generated_outputs
@@ -210,8 +214,13 @@ def _toxicity_local(generated_outputs: List[str]) -> List[float]:
     global _toxicity_model_path
     # this pipeline output predict probability for each text on each label.
     # the output format is List[List[Dict(str)]]
+    from langcheck.metrics import _model_manager
+    tokenizer, model = _model_manager.fetch_model(language='zh',
+                                                  metric_type="toxicity")
+
     _toxicity_pipeline = pipeline('text-classification',
-                                  model=_toxicity_model_path,
+                                  model=model,
+                                  tokenizer=tokenizer,  # type: ignore[reportOptionalIterable]  # NOQA: E501
                                   top_k=5)
 
     # {'Normal': 0, 'Pulp': 1, 'Sex': 2, 'Other Risk': 3, 'Adult': 4}

src/langcheck/metrics/zh/source_based_text_quality.py

@@ -86,8 +86,10 @@ def factual_consistency(
 
     global _factual_consistency_translation_pipeline
     if _factual_consistency_translation_pipeline is None:
+        from langcheck.metrics import _model_manager
+        tokenizer, model = _model_manager.fetch_model(language='zh', metric_type='factual')   # NOQA: E501
         _factual_consistency_translation_pipeline = pipeline(
-            'translation', model=_factual_consistency_translation_model_path)
+            'translation', model=model, tokenizer=tokenizer)  # type: ignore
 
     # Translate the sources and generated outputs to English.
     # Currently, the type checks are not working for the pipeline, since