Another multilingual approach

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.0.27
+current_version = 0.0.28-dev1
 tag = False
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(\-(?P<release>[a-z]+)(?P<build>\d+))?
 serialize = 

langkit/__init__.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Dict, List
+from typing import Dict, List, Optional, Set
 from .extract import extract
 import importlib.resources as resources
 
@@ -14,11 +14,18 @@ class LangKitConfig:
     pattern_file_path: str = field(
         default_factory=lambda: _resource_filename("pattern_groups.json")
     )
+    response_pattern_file_path: Optional[str] = field(
+        default_factory=lambda: _resource_filename("pattern_groups.json")
+    )
     metric_name_map: Dict[str, str] = field(default_factory=dict)
     theme_file_path: str = field(
         default_factory=lambda: _resource_filename("themes.json")
     )
-    transformer_name: str = "sentence-transformers/all-MiniLM-L6-v2"
+    response_theme_file_path: str = field(
+        default_factory=lambda: _resource_filename("themes.json")
+    )
+    transformer_name: Optional[str] = "sentence-transformers/all-MiniLM-L6-v2"
+    response_transformer_name: Optional[str] = "sentence-transformers/all-MiniLM-L6-v2"
     topics: List[str] = field(
         default_factory=lambda: [
             "law",
@@ -29,30 +36,122 @@ class LangKitConfig:
             "support",
         ]
     )
+    response_topics: List[str] = field(
+        default_factory=lambda: [
+            "law",
+            "finance",
+            "medical",
+            "education",
+            "politics",
+            "support",
+        ]
+    )
     nlp_scores: list = field(
         default_factory=lambda: [
             "bleu",
             "rouge",
             "meteor",
         ]
     )
-    reference_corpus: str = ""
+    reference_corpus: Optional[str] = ""
     injections_base_url = (
         "https://whylabs-public.s3.us-west-2.amazonaws.com/langkit/data/injections/"
     )
     data_folder: str = "langkit_data"
     rouge_type: str = "rouge1"
-    sentiment_lexicon: str = "vader_lexicon"
-    topic_model_path: str = "MoritzLaurer/mDeBERTa-v3-base-xnli-multilingual-nli-2mil7"
-    topic_classifier: str = "zero-shot-classification"
-    toxicity_model_path: str = "martin-ha/toxic-comment-model"
+    sentiment_lexicon: Optional[str] = "vader_lexicon"
+    response_sentiment_lexicon: Optional[str] = "vader_lexicon"
+    topic_model_path: Optional[
+        str
+    ] = "MoritzLaurer/mDeBERTa-v3-base-xnli-multilingual-nli-2mil7"
+    response_topic_model_path: Optional[
+        str
+    ] = "MoritzLaurer/mDeBERTa-v3-base-xnli-multilingual-nli-2mil7"
+    topic_classifier: Optional[str] = "zero-shot-classification"
+    response_topic_classifier: Optional[str] = "zero-shot-classification"
+    toxicity_model_path: Optional[str] = "martin-ha/toxic-comment-model"
+    response_toxicity_model_path: Optional[str] = "martin-ha/toxic-comment-model"
+    injections_transformer_name: Optional[str] = "all-MiniLM-L6-v2"
+    injections_version: Optional[str] = "v1"
+    prompt_languages: Optional[Set[str]] = field(default_factory=lambda: {"en"})
+    response_languages: Optional[Set[str]] = field(default_factory=lambda: {"en"})
+    sentiment_model_path: Optional[
+        str
+    ] = "lxyuan/distilbert-base-multilingual-cased-sentiments-student"
+    response_sentiment_model_path: Optional[
+        str
+    ] = "lxyuan/distilbert-base-multilingual-cased-sentiments-student"
 
 
 prompt_column: str = "prompt"
 response_column: str = "response"
 lang_config = LangKitConfig()
 
 
+# Override default models/parameters per language
+multi_lang_config: Dict[Optional[str], LangKitConfig] = {
+    None: LangKitConfig(),
+    "": LangKitConfig(),
+    "ar": LangKitConfig(
+        prompt_languages={"ar"},
+        response_languages={"ar"},
+        injections_transformer_name=None,
+        reference_corpus=None,
+        sentiment_lexicon=None,
+        response_sentiment_lexicon=None,
+        topic_model_path=None,
+        response_topic_model_path=None,
+        toxicity_model_path=None,
+        response_toxicity_model_path=None,
+        transformer_name=None,
+        response_transformer_name=None,
+    ),
+    "en": LangKitConfig(),
+    "es": LangKitConfig(
+        prompt_languages={"es"},
+        response_languages={"es"},
+        injections_transformer_name=None,
+        reference_corpus=None,
+        sentiment_lexicon=None,
+        response_sentiment_lexicon=None,
+        topic_model_path=None,
+        response_topic_model_path=None,
+        toxicity_model_path=None,
+        response_toxicity_model_path=None,
+        transformer_name=None,
+        response_transformer_name=None,
+    ),
+    "it": LangKitConfig(
+        prompt_languages={"it"},
+        response_languages={"it"},
+        injections_transformer_name=None,
+        reference_corpus=None,
+        sentiment_lexicon=None,
+        response_sentiment_lexicon=None,
+        topic_model_path=None,
+        response_topic_model_path=None,
+        toxicity_model_path=None,
+        response_toxicity_model_path=None,
+        transformer_name=None,
+        response_transformer_name=None,
+    ),
+    "pt": LangKitConfig(
+        prompt_languages={"pt"},
+        response_languages={"pt"},
+        injections_transformer_name=None,
+        reference_corpus=None,
+        sentiment_lexicon=None,
+        response_sentiment_lexicon=None,
+        topic_model_path=None,
+        response_topic_model_path=None,
+        toxicity_model_path="dougtrajano/toxicity-type-detection",
+        response_toxicity_model_path="dougtrajano/toxicity-type-detection",
+        transformer_name=None,
+        response_transformer_name=None,
+    ),
+}
+
+
 def package_version(package: str = __package__) -> str:
     """Calculate version number based on pyproject.toml"""
     try:

langkit/all_metrics.py

@@ -4,7 +4,7 @@
 
 from langkit.metadata import attach_schema_metadata
 
-from langkit import LangKitConfig
+from langkit import LangKitConfig, multi_lang_config
 from langkit import injections
 from langkit import topics
 from langkit import regexes
@@ -15,14 +15,36 @@
 from langkit import input_output
 
 
-def init(config: Optional[LangKitConfig] = None) -> DeclarativeSchema:
-    injections.init(config=config)
-    topics.init(config=config)
-    regexes.init(config=config)
-    sentiment.init(config=config)
-    textstat.init(config=config)
-    themes.init(config=config)
-    toxicity.init(config=config)
-    input_output.init(config=config)
-    text_schema = attach_schema_metadata(udf_schema(), "all_metrics")
+def init(
+    language: Optional[str] = None,
+    config: Optional[LangKitConfig] = None,
+    schema_name: str = "",
+) -> DeclarativeSchema:
+    injections.init(
+        language, config=config or multi_lang_config[language], schema_name=schema_name
+    )
+    topics.init(
+        language, config=config or multi_lang_config[language], schema_name=schema_name
+    )
+    regexes.init(
+        language, config=config or multi_lang_config[language], schema_name=schema_name
+    )
+    sentiment.init(
+        language, config=config or multi_lang_config[language], schema_name=schema_name
+    )
+    textstat.init(
+        language, config=config or multi_lang_config[language], schema_name=schema_name
+    )
+    themes.init(
+        language, config=config or multi_lang_config[language], schema_name=schema_name
+    )
+    toxicity.init(
+        language, config=config or multi_lang_config[language], schema_name=schema_name
+    )
+    input_output.init(
+        language, config=config or multi_lang_config[language], schema_name=schema_name
+    )
+    text_schema = attach_schema_metadata(
+        udf_schema(schema_name=schema_name), "all_metrics"
+    )
     return text_schema

langkit/count_regexes.py

@@ -1,3 +1,4 @@
+from collections import defaultdict
 from copy import deepcopy
 from logging import getLogger
 
@@ -6,14 +7,20 @@
 from langkit import LangKitConfig, lang_config, prompt_column, response_column
 from whylogs.core.stubs import pd
 from typing import Dict, List, Optional, Set, Union
+from langkit.whylogs.unreg import unregister_udfs  # replace with whylogs 1.3.12
 
 diagnostic_logger = getLogger(__name__)
 
 
 pattern_loader = PatternLoader()
+response_pattern_loader = PatternLoader()
+
+_initialized = False
 
 
 def count_patterns(group, text: str) -> int:
+    if not _initialized:
+        init()
     count = 0
     for expression in group["expressions"]:
         if expression.search(text):
@@ -29,48 +36,59 @@ def wrappee(text: Union[pd.DataFrame, Dict[str, List]]) -> Union[pd.Series, List
     return wrappee
 
 
-_registered: Set[str] = set()
-
-
-def _unregister():
-    # WARNING: Uses private whylogs internals. Do not copy this code.
-    # TODO: Add proper whylogs API to support this.
-    from whylogs.experimental.core.udf_schema import _multicolumn_udfs
+_registered: Dict[str, Set[str]] = defaultdict(
+    set
+)  # _registered[schema_name] -> set of registered UDF names
 
-    global _multicolumn_udfs, _registered
-    _multicolumn_udfs[""] = [
-        u for u in _multicolumn_udfs[""] if list(u.udfs.keys())[0] not in _registered
-    ]
-    _registered = set()
 
-
-def _register_udfs():
+def _register_udfs(language: str, schema_name: str):
     global _registered
-    _unregister()
+    unregister_udfs(_registered[schema_name], language, schema_name)
+    _registered[schema_name] = set()
+
     regex_groups = pattern_loader.get_regex_groups()
     if regex_groups is not None:
-        for column in [prompt_column, response_column]:
-            for group in regex_groups:
-                udf_name = f"{column}.{group['name']}_count"
-                register_dataset_udf(
-                    [column],
-                    udf_name=udf_name,
-                )(wrapper(group, column))
-                _registered.add(udf_name)
+        column = prompt_column
+        for group in regex_groups:
+            udf_name = f"{column}.{group['name']}_count"
+            register_dataset_udf(
+                [column],
+                udf_name=udf_name,
+                namespace=language,
+                schema_name=schema_name,
+            )(wrapper(group, column))
+            _registered[schema_name].add(udf_name)
+
+    regex_groups = response_pattern_loader.get_regex_groups()
+    if regex_groups is not None:
+        column = response_column
+        for group in regex_groups:
+            udf_name = f"{column}.{group['name']}_count"
+            register_dataset_udf(
+                [column],
+                udf_name=udf_name,
+                namespace=language,
+                schema_name=schema_name,
+            )(wrapper(group, column))
+            _registered[schema_name].add(udf_name)
 
 
 def init(
-    pattern_file_path: Optional[str] = None, config: Optional[LangKitConfig] = None
+    language: Optional[str] = None,
+    pattern_file_path: Optional[str] = None,
+    config: Optional[LangKitConfig] = None,
+    response_pattern_file_path: Optional[str] = None,
+    schema_name: str = "",
 ):
+    global _initialized
+    _initialized = True
+    language = language or ""
     config = deepcopy(config or lang_config)
     if pattern_file_path:
         config.pattern_file_path = pattern_file_path
-
-    global pattern_loader
-    pattern_loader = PatternLoader(config)
-    pattern_loader.update_patterns()
-
-    _register_udfs()
-
-
-init()
+    if response_pattern_file_path:
+        config.response_pattern_file_path = response_pattern_file_path
+    global pattern_loader, response_pattern_loader
+    pattern_loader = PatternLoader(config.pattern_file_path)
+    response_pattern_loader = PatternLoader(config.response_pattern_file_path)
+    _register_udfs(language, schema_name)