add nlp score metrics (#87)

* add nlp score metrics

---------

Co-authored-by: Felipe Adachi <[email protected]>

langkit/__init__.py

@@ -22,6 +22,14 @@ class LangKitConfig:
             "support",
         ]
     )
+    nlp_scores: list = field(
+        default_factory=lambda: [
+            "bleu",
+            "rouge",
+            "meteor",
+        ]
+    )
+    reference_corpus: str = ""
 
 
 def package_version(package: str = __package__) -> str:

langkit/nlp_scores.py

@@ -0,0 +1,66 @@
+from whylogs.experimental.core.metrics.udf_metric import register_metric_udf
+import evaluate
+from . import LangKitConfig
+from logging import getLogger
+
+lang_config = LangKitConfig()
+_corpus = lang_config.reference_corpus
+_scores = lang_config.nlp_scores
+_rouge_type = "rouge1"
+
+diagnostic_logger = getLogger(__name__)
+
+
+def register_score_udfs():
+    if _corpus:
+        for score in _scores:
+            if "bleu" in score:
+                bleu = evaluate.load("bleu")
+
+                @register_metric_udf(col_name=lang_config.response_column)
+                def bleu_score(text: str) -> float:
+                    return bleu.compute(predictions=[text], references=[_corpus])[
+                        "bleu"
+                    ]
+
+            if "rouge" in score:
+                rouge = evaluate.load("rouge")
+
+                @register_metric_udf(col_name=lang_config.response_column)
+                def rouge_score(text: str) -> float:
+                    return rouge.compute(
+                        predictions=[text],
+                        references=[_corpus],
+                        rouge_types=[_rouge_type],
+                    )[_rouge_type]
+
+            if "meteor" in score:
+                meteor = evaluate.load("meteor")
+
+                @register_metric_udf(col_name=lang_config.response_column)
+                def meteor_score(text: str) -> float:
+                    return meteor.compute(predictions=[text], references=[_corpus])[
+                        "meteor"
+                    ]
+
+    else:
+        diagnostic_logger.warning(
+            "No reference corpus provided for NLP scores. Skipping NLP scores."
+        )
+
+
+def init(corpus=None, scores=[], rouge_type=None):
+    global _corpus
+    global _scores
+    global _rouge_type
+    if corpus:
+        _corpus = corpus
+    if scores:
+        _scores = scores
+    if rouge_type:
+        _rouge_type = rouge_type
+
+    register_score_udfs()
+
+
+init()

langkit/tests/__init__.py

@@ -22,6 +22,14 @@ class LangKitConfig:
             "support",
         ]
     )
+    nlp_scores: list = field(
+        default_factory=lambda: [
+            "bleu",
+            "rouge",
+            "meteor",
+        ]
+    )
+    reference_corpus: str = ""
 
 
 def package_version(package: str = __package__) -> str:

langkit/tests/test_nlp_scores.py

@@ -0,0 +1,23 @@
+import whylogs as why
+import pytest
+
+
+@pytest.mark.load
+def test_bleu_score():
+    from langkit import nlp_scores  # noqa
+    from whylogs.experimental.core.udf_schema import udf_schema
+
+    nlp_scores.init(
+        scores=["bleu"], corpus="The quick brown fox jumps over the lazy dog"
+    )
+    text_schema = udf_schema()
+    profile = why.log(
+        {"response": "The quick dog jumps over the lazy brown fox"}, schema=text_schema
+    ).profile()
+    max_score = (
+        profile.view()
+        .get_column("response")
+        .get_metrics()[-1]
+        .to_summary_dict()["bleu_score:distribution/max"]
+    )
+    assert round(max_score, 2) == 0.42

pyproject.toml

@@ -20,6 +20,7 @@ datasets = {version ="^2.12.0", optional = true}
 openai = {version ="^0.27.6", optional = true}
 nltk = {version ="^3.8.1", optional = true}
 sentence-transformers = {version ="^2.2.2", optional = true}
+evaluate = {version = "^0.4.0", optional = true}
 
 
 [tool.poetry.group.dev.dependencies]
@@ -40,6 +41,7 @@ all = [
     "openai",
     "nltk",
     "sentence-transformers",
+    "evaluate",
 ]