From 3c0e2a79e21904a89da0466ea5ccd18e12d40268 Mon Sep 17 00:00:00 2001 From: Richard Rogers Date: Mon, 30 Oct 2023 03:41:29 -0700 Subject: [PATCH 1/3] Another multilingual approach --- .bumpversion.cfg | 2 +- langkit/__init__.py | 113 ++++++- langkit/all_metrics.py | 22 +- langkit/count_regexes.py | 74 +++-- langkit/examples/Multilingual3_Example.ipynb | 303 +++++++++++++++++++ langkit/injections.py | 60 ++-- langkit/input_output.py | 64 ++-- langkit/light_metrics.py | 10 +- langkit/llm_metrics.py | 18 +- langkit/nlp_scores.py | 32 +- langkit/pattern_loader.py | 17 +- langkit/regexes.py | 103 +++---- langkit/response_hallucination.py | 47 +-- langkit/sentiment.py | 161 ++++++++-- langkit/tests/test_injections.py | 1 + langkit/tests/test_textstat.py | 14 +- langkit/tests/test_themes.py | 22 +- langkit/tests/test_toxicity.py | 3 + langkit/tests/test_translation.py | 67 ++++ langkit/textstat.py | 76 +++-- langkit/themes.py | 142 ++++++--- langkit/topics.py | 68 ++++- langkit/toxicity.py | 96 ++++-- langkit/translator.py | 48 +++ langkit/whylogs/unreg.py | 35 +++ pyproject.toml | 2 +- 26 files changed, 1257 insertions(+), 343 deletions(-) create mode 100644 langkit/examples/Multilingual3_Example.ipynb create mode 100644 langkit/tests/test_translation.py create mode 100644 langkit/translator.py create mode 100644 langkit/whylogs/unreg.py diff --git a/.bumpversion.cfg b/.bumpversion.cfg index 4c6411c0..685fdba0 100644 --- a/.bumpversion.cfg +++ b/.bumpversion.cfg @@ -1,5 +1,5 @@ [bumpversion] -current_version = 0.0.25 +current_version = 0.0.26-dev1 tag = False parse = (?P\d+)\.(?P\d+)\.(?P\d+)(\-(?P[a-z]+)(?P\d+))? serialize = diff --git a/langkit/__init__.py b/langkit/__init__.py index a0c131ba..f8ecc1d6 100644 --- a/langkit/__init__.py +++ b/langkit/__init__.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, field -from typing import Dict, List +from typing import Dict, List, Optional, Set 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,6 +36,16 @@ 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", @@ -36,16 +53,34 @@ class LangKitConfig: "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" @@ -53,6 +88,70 @@ class LangKitConfig: 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: diff --git a/langkit/all_metrics.py b/langkit/all_metrics.py index 3d66235d..6e280c55 100644 --- a/langkit/all_metrics.py +++ b/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,16 @@ 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) +def init( + language: Optional[str] = None, config: Optional[LangKitConfig] = None +) -> DeclarativeSchema: + injections.init(language, config=config or multi_lang_config[language]) + topics.init(language, config=config or multi_lang_config[language]) + regexes.init(language, config=config or multi_lang_config[language]) + sentiment.init(language, config=config or multi_lang_config[language]) + textstat.init(language, config=config or multi_lang_config[language]) + themes.init(language, config=config or multi_lang_config[language]) + toxicity.init(language, config=config or multi_lang_config[language]) + input_output.init(language, config=config or multi_lang_config[language]) text_schema = attach_schema_metadata(udf_schema(), "all_metrics") return text_schema diff --git a/langkit/count_regexes.py b/langkit/count_regexes.py index 3e9dc871..fa12bf26 100644 --- a/langkit/count_regexes.py +++ b/langkit/count_regexes.py @@ -6,14 +6,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): @@ -32,45 +38,49 @@ def wrappee(text: Union[pd.DataFrame, Dict[str, List]]) -> Union[pd.Series, List _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 - - 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): global _registered - _unregister() + unregister_udfs(_registered) 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, + schema_name=language, + )(wrapper(group, column)) + _registered.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, + schema_name=language, + )(wrapper(group, column)) + _registered.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, ): + 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) diff --git a/langkit/examples/Multilingual3_Example.ipynb b/langkit/examples/Multilingual3_Example.ipynb new file mode 100644 index 00000000..f32fe675 --- /dev/null +++ b/langkit/examples/Multilingual3_Example.ipynb @@ -0,0 +1,303 @@ +{ + "nbformat": 4, + "nbformat_minor": 0, + "metadata": { + "colab": { + "provenance": [] + }, + "kernelspec": { + "name": "python3", + "display_name": "Python 3" + }, + "language_info": { + "name": "python" + } + }, + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "ZDUtcLEmtIlk" + }, + "outputs": [], + "source": [ + "%pip install langkit[all]==0.0.24.dev2" + ] + }, + { + "cell_type": "markdown", + "source": [ + "Let's say we want the standard LLM metrics for English prompts, but we don't want any metrics on the responses since they're in a different unsupported language. We can use the default configurations for the prompt metrics and set the response metric configurations to `None` to turn them off.\n" + ], + "metadata": { + "id": "mqPATL2OvFAG" + } + }, + { + "cell_type": "code", + "source": [ + "import langkit.llm_metrics as llmm\n", + "from langkit import LangKitConfig\n", + "\n", + "no_response_config = LangKitConfig(\n", + " prompt_languages={\"en\"},\n", + " response_languages={\"pt\"},\n", + " response_pattern_file_path=None, # regexes are semi-international, could leave this as default\n", + " sentiment_lexicon=None, # use multilingual sentiment model for the prompt\n", + " response_sentiment_lexicon=None,\n", + " response_sentiment_model_path=None,\n", + " response_theme_file_path=None,\n", + " response_topic_model_path=None,\n", + " response_toxicity_model_path=None,\n", + " response_transformer_name=None,\n", + ")\n", + "schema = llmm.init(config=no_response_config)" + ], + "metadata": { + "id": "0U_TM-3Xtx-Z" + }, + "execution_count": 2, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "import whylogs as why\n", + "from langkit.whylogs.samples import load_chats, show_first_chat\n", + "\n", + "chats = load_chats()\n", + "results = why.log(chats, schema=schema)\n", + "\n", + "for column in results.view().get_columns().keys():\n", + " print(column)" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "dQDvPLvR8E_U", + "outputId": "c4450822-3795-40a7-fb03-36109d35d745" + }, + "execution_count": 3, + "outputs": [ + { + "output_type": "stream", + "name": "stderr", + "text": [ + "/usr/local/lib/python3.10/dist-packages/transformers/pipelines/text_classification.py:105: UserWarning: `return_all_scores` is now deprecated, if want a similar functionality use `top_k=None` instead of `return_all_scores=True` or `top_k=1` instead of `return_all_scores=False`.\n", + " warnings.warn(\n" + ] + }, + { + "output_type": "stream", + "name": "stdout", + "text": [ + "⚠️ No session found. Call whylogs.init() to initialize a session and authenticate. See https://docs.whylabs.ai/docs/whylabs-whylogs-init for more information.\n", + "prompt\n", + "response\n", + "prompt.has_patterns\n", + "prompt.sentiment_multi\n", + "prompt.flesch_reading_ease\n", + "prompt.automated_readability_index\n", + "prompt.aggregate_reading_level\n", + "prompt.syllable_count\n", + "prompt.lexicon_count\n", + "prompt.sentence_count\n", + "prompt.character_count\n", + "prompt.letter_count\n", + "prompt.polysyllable_count\n", + "prompt.monosyllable_count\n", + "prompt.difficult_words\n", + "prompt.jailbreak_similarity\n", + "prompt.toxicity\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "source": [ + "We have a multilingual sentiment model that includes Portuguese, so we can enable it to track response sentiment." + ], + "metadata": { + "id": "bH4ZWQBLe0hS" + } + }, + { + "cell_type": "code", + "source": [ + "pt_sentiment_config = LangKitConfig(\n", + " prompt_languages={\"en\"},\n", + " response_languages={\"pt\"},\n", + " response_pattern_file_path=None,\n", + " sentiment_lexicon=None, # disable English-only sentiment model\n", + " response_sentiment_lexicon=None,\n", + " #response_sentiment_model_path=None, # enable (don't disable) multilingual sentiment model\n", + " response_theme_file_path=None,\n", + " response_topic_model_path=None,\n", + " response_toxicity_model_path=None,\n", + " response_transformer_name=None,\n", + ")\n", + "schema = llmm.init(config=pt_sentiment_config)\n", + "results = why.log(chats, schema=schema)\n", + "\n", + "for column in results.view().get_columns().keys():\n", + " print(column)\n" + ], + "metadata": { + "colab": { + "base_uri": "https://localhost:8080/" + }, + "id": "HxRf98Ujfr6D", + "outputId": "73e9efc0-65d6-41e8-b225-e4296a228966" + }, + "execution_count": 4, + "outputs": [ + { + "output_type": "stream", + "name": "stderr", + "text": [ + "/usr/local/lib/python3.10/dist-packages/transformers/pipelines/text_classification.py:105: UserWarning: `return_all_scores` is now deprecated, if want a similar functionality use `top_k=None` instead of `return_all_scores=True` or `top_k=1` instead of `return_all_scores=False`.\n", + " warnings.warn(\n", + "/usr/local/lib/python3.10/dist-packages/transformers/pipelines/text_classification.py:105: UserWarning: `return_all_scores` is now deprecated, if want a similar functionality use `top_k=None` instead of `return_all_scores=True` or `top_k=1` instead of `return_all_scores=False`.\n", + " warnings.warn(\n" + ] + }, + { + "output_type": "stream", + "name": "stdout", + "text": [ + "prompt\n", + "response\n", + "prompt.has_patterns\n", + "prompt.sentiment_multi\n", + "response.sentiment_multi\n", + "prompt.flesch_reading_ease\n", + "prompt.automated_readability_index\n", + "prompt.aggregate_reading_level\n", + "prompt.syllable_count\n", + "prompt.lexicon_count\n", + "prompt.sentence_count\n", + "prompt.character_count\n", + "prompt.letter_count\n", + "prompt.polysyllable_count\n", + "prompt.monosyllable_count\n", + "prompt.difficult_words\n", + "prompt.jailbreak_similarity\n", + "prompt.toxicity\n" + ] + } + ] + }, + { + "cell_type": "markdown", + "source": [ + "The toxicity module supports translation, so we can configure a Portuguese to English translator and track response toxicity with the English-only model." + ], + "metadata": { + "id": "jnpjxIGKg9ZV" + } + }, + { + "cell_type": "code", + "source": [ + "\"\"\"\n", + "from transformers import AutoProcessor, SeamlessM4TModel\n", + "from langkit.translator import Translator\n", + "\n", + "class Seamless_M4T(Translator):\n", + " def __init__(self, src_lang=\"eng\", tgt_lang=\"eng\"):\n", + " self._processor = AutoProcessor.from_pretrained(\"facebook/hf-seamless-m4t-medium\")\n", + " self._model = SeamlessM4TModel.from_pretrained(\"facebook/hf-seamless-m4t-medium\")\n", + " self._src_lang = src_lang\n", + " self._tgt_lang = tgt_lang\n", + "\n", + " def translate(self, text):\n", + " text_inputs = self._processor(text = text, src_lang=self._src_lang, tgt_lang=self._tgt_lang, return_tensors=\"pt\")\n", + " output_tokens = self._model.generate(**text_inputs, tgt_lang=self._tgt_lang, generate_speech=False)\n", + " return self._processor.decode(output_tokens[0].tolist()[0], skip_special_tokens=True)\n", + "\n", + "print(Seamless_M4T(\"por\", \"eng\").translate(\"Olá, o meu cão é lindo.\"))\n", + "\"\"\"\n", + "\n", + "from langkit.translator import Translator\n", + "class FakeTranslator(Translator):\n", + " def translate(self, text):\n", + " return \"Hello, my dog is beautiful\"" + ], + "metadata": { + "id": "b5quNo4sg-Mb" + }, + "execution_count": 5, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "import langkit.toxicity as tx\n", + "tx.RESPONSE_TRANSLATOR = FakeTranslator() # Seamless_M4T(\"por\", \"eng\")\n", + "\n", + "translated_config = LangKitConfig(\n", + " prompt_languages={\"en\"},\n", + " response_languages={\"pt\"},\n", + " response_pattern_file_path=None,\n", + " sentiment_lexicon=None,\n", + " response_sentiment_lexicon=None,\n", + " response_sentiment_model_path=None,\n", + " response_theme_file_path=None,\n", + " response_topic_model_path=None,\n", + " # response_toxicity_model_path=None,\n", + " response_transformer_name=None,\n", + ")\n", + "schema = llmm.init(config=translated_config)\n", + "results = why.log({\"prompt\": \"Hi, how's your dog?\", \"response\": \"Olá, o meu cão é lindo.\"}, schema=schema)\n", + "\n", + "for column in results.view().get_columns().keys():\n", + " print(column)" + ], + "metadata": { + "id": "YNo4hIxzhxBq", + "outputId": "e4a2669e-1e55-4a7b-9aef-931391fc2656", + "colab": { + "base_uri": "https://localhost:8080/" + } + }, + "execution_count": 6, + "outputs": [ + { + "output_type": "stream", + "name": "stderr", + "text": [ + "/usr/local/lib/python3.10/dist-packages/transformers/pipelines/text_classification.py:105: UserWarning: `return_all_scores` is now deprecated, if want a similar functionality use `top_k=None` instead of `return_all_scores=True` or `top_k=1` instead of `return_all_scores=False`.\n", + " warnings.warn(\n" + ] + }, + { + "output_type": "stream", + "name": "stdout", + "text": [ + "prompt\n", + "response\n", + "prompt.has_patterns\n", + "prompt.sentiment_multi\n", + "prompt.flesch_reading_ease\n", + "prompt.automated_readability_index\n", + "prompt.aggregate_reading_level\n", + "prompt.syllable_count\n", + "prompt.lexicon_count\n", + "prompt.sentence_count\n", + "prompt.character_count\n", + "prompt.letter_count\n", + "prompt.polysyllable_count\n", + "prompt.monosyllable_count\n", + "prompt.difficult_words\n", + "prompt.jailbreak_similarity\n", + "prompt.toxicity\n", + "response.toxicity\n" + ] + } + ] + } + ] +} diff --git a/langkit/injections.py b/langkit/injections.py index 17941d22..a38728f8 100644 --- a/langkit/injections.py +++ b/langkit/injections.py @@ -1,8 +1,9 @@ from copy import deepcopy -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional, Set, Union from whylogs.core.stubs import pd from whylogs.experimental.core.udf_schema import register_dataset_udf from langkit import LangKitConfig, lang_config, prompt_column +from langkit.whylogs.unreg import unregister_udfs from sentence_transformers import SentenceTransformer import requests from io import BytesIO @@ -12,10 +13,26 @@ import os import torch -_prompt = prompt_column + _index_embeddings = None _transformer_model = None +_initialized = False + +def injection(prompt: Union[Dict[str, List], pd.DataFrame]) -> Union[List, pd.Series]: + if not _initialized: + init() + global _transformer_model + global _index_embeddings + if _transformer_model is None: + raise ValueError("Injections - transformer model not initialized") + embeddings = _transformer_model.encode(prompt[prompt_column]) + faiss.normalize_L2(embeddings) + if _index_embeddings is None: + raise ValueError("Injections - index embeddings not initialized") + dists, _ = _index_embeddings.search(x=embeddings, k=1) + return dists.flatten().tolist() + _USE_CUDA = torch.cuda.is_available() and not bool( os.environ.get("LANGKIT_NO_CUDA", False) @@ -30,18 +47,29 @@ def download_embeddings(url): return array +_registered: Set[str] = set() + + def init( + language: Optional[str] = None, transformer_name: Optional[str] = None, version: Optional[str] = None, config: Optional[LangKitConfig] = None, ): + global _initialized + _initialized = True + global _registered + unregister_udfs(_registered) config = config or deepcopy(lang_config) global _transformer_model global _index_embeddings - if not transformer_name: - transformer_name = "all-MiniLM-L6-v2" - if not version: - version = "v1" + transformer_name = transformer_name or config.injections_transformer_name + version = version or config.injections_version + + if transformer_name is None or version is None: + _transformer_model = None + return + _transformer_model = SentenceTransformer(transformer_name, device=_device) path = f"index_embeddings_{transformer_name}_harm_{version}.npy" @@ -79,20 +107,6 @@ def init( raise ValueError( f"Injections - unable to deserialize index to {embeddings_path}. Error: {deserialization_error}" ) - - -@register_dataset_udf([_prompt], f"{_prompt}.injection") -def injection(prompt: Union[Dict[str, List], pd.DataFrame]) -> Union[List, pd.Series]: - global _transformer_model - global _index_embeddings - if _transformer_model is None: - raise ValueError("Injections - transformer model not initialized") - embeddings = _transformer_model.encode(prompt[_prompt]) - faiss.normalize_L2(embeddings) - if _index_embeddings is None: - raise ValueError("Injections - index embeddings not initialized") - dists, _ = _index_embeddings.search(x=embeddings, k=1) - return dists.flatten().tolist() - - -init() + if _index_embeddings and _transformer_model: + register_dataset_udf([prompt_column], f"{prompt_column}.injection")(injection) + _registered.add(f"{prompt_column}.injection") diff --git a/langkit/input_output.py b/langkit/input_output.py index 434cb298..8a580263 100644 --- a/langkit/input_output.py +++ b/langkit/input_output.py @@ -1,11 +1,12 @@ from copy import deepcopy from logging import getLogger -from typing import Callable, Optional +from typing import Callable, Optional, Set from sentence_transformers import util from whylogs.experimental.core.udf_schema import register_dataset_udf from langkit import LangKitConfig, lang_config, prompt_column, response_column from langkit.transformer import Encoder +from langkit.whylogs.unreg import unregister_udfs _prompt = prompt_column _response = response_column @@ -13,26 +14,14 @@ _transformer_model = None -diagnostic_logger = getLogger(__name__) - +_initialized = False -def init( - transformer_name: Optional[str] = None, - custom_encoder: Optional[Callable] = None, - config: Optional[LangKitConfig] = None, -): - config = config or deepcopy(lang_config) - global _transformer_model - if transformer_name is None and custom_encoder is None: - transformer_name = config.transformer_name - _transformer_model = Encoder(transformer_name, custom_encoder) - - -init() +diagnostic_logger = getLogger(__name__) -@register_dataset_udf([_prompt, _response], f"{_response}.relevance_to_{_prompt}") def prompt_response_similarity(text): + if not _initialized: + init() global _transformer_model if _transformer_model is None: @@ -53,3 +42,44 @@ def prompt_response_similarity(text): ) series_result.append(None) return series_result + + +_registered: Set[str] = set() + + +def init( + language: Optional[str] = None, + transformer_name: Optional[str] = None, + custom_encoder: Optional[Callable] = None, + config: Optional[LangKitConfig] = None, +): + global _initialized + _initialized = True + global _registered + unregister_udfs(_registered) + if transformer_name and custom_encoder: + raise ValueError( + "Only one of transformer_name or encoder can be specified, not both." + ) + config = config or deepcopy(lang_config) + global _transformer_model + response_transformer_name = ( + transformer_name or config.response_transformer_name + ) # not a bug :) + transformer_name = transformer_name or config.transformer_name + + if transformer_name != response_transformer_name: # can't evaluate across langauges + _transformer_model = None + return + + if transformer_name is None and custom_encoder is None: # metric turned off + _transformer_model = None + return + + transformer_name = None if custom_encoder else transformer_name + _transformer_model = Encoder(transformer_name, custom_encoder) + register_dataset_udf( + [prompt_column, response_column], + f"{response_column}.relevance_to_{prompt_column}", + )(prompt_response_similarity) + _registered.add(f"{response_column}.relevance_to_{prompt_column}") diff --git a/langkit/light_metrics.py b/langkit/light_metrics.py index 99f884ae..c8750165 100644 --- a/langkit/light_metrics.py +++ b/langkit/light_metrics.py @@ -2,15 +2,17 @@ from whylogs.experimental.core.udf_schema import udf_schema from whylogs.core.schema import DeclarativeSchema -from langkit import LangKitConfig +from langkit import LangKitConfig, multi_lang_config from langkit.metadata import attach_schema_metadata from langkit import regexes from langkit import textstat -def init(config: Optional[LangKitConfig] = None) -> DeclarativeSchema: - regexes.init(config=config) - textstat.init(config=config) +def init( + language: Optional[str] = None, config: Optional[LangKitConfig] = None +) -> DeclarativeSchema: + regexes.init(language, config=config or multi_lang_config[language]) + textstat.init(language, config=config or multi_lang_config[language]) text_schema = attach_schema_metadata(udf_schema(), "light_metrics") return text_schema diff --git a/langkit/llm_metrics.py b/langkit/llm_metrics.py index 1019d1b5..062e79e1 100644 --- a/langkit/llm_metrics.py +++ b/langkit/llm_metrics.py @@ -1,5 +1,5 @@ from langkit.metadata import attach_schema_metadata -from langkit import LangKitConfig +from langkit import LangKitConfig, multi_lang_config from logging import getLogger from typing import Optional from whylogs.experimental.core.udf_schema import udf_schema @@ -20,13 +20,15 @@ ) -def init(config: Optional[LangKitConfig] = None) -> DeclarativeSchema: - 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) +def init( + language: Optional[str] = None, config: Optional[LangKitConfig] = None +) -> DeclarativeSchema: + regexes.init(language, config=config or multi_lang_config[language]) + sentiment.init(language, config=config or multi_lang_config[language]) + textstat.init(language, config=config or multi_lang_config[language]) + themes.init(language, config=config or multi_lang_config[language]) + toxicity.init(language, config=config or multi_lang_config[language]) + input_output.init(language, config=config or multi_lang_config[language]) text_schema = attach_schema_metadata(udf_schema(), "llm_metrics") return text_schema diff --git a/langkit/nlp_scores.py b/langkit/nlp_scores.py index 1103d536..e62397c6 100644 --- a/langkit/nlp_scores.py +++ b/langkit/nlp_scores.py @@ -4,29 +4,31 @@ import evaluate from langkit import LangKitConfig, lang_config, response_column from logging import getLogger +from langkit.whylogs.unreg import unregister_udfs -_corpus: str = lang_config.reference_corpus +_corpus: Optional[str] = lang_config.reference_corpus _scores: List[str] = lang_config.nlp_scores _rouge_type: str = lang_config.rouge_type diagnostic_logger = getLogger(__name__) +_initialized = False -_bleu_registered = False -_rouge_registered = False -_meteor_registered = False +_registered: Set[str] = set() def _register_score_udfs(): - global _bleu_registered, _rouge_registered, _meteor_registered - + if not _initialized: + init() + global _registered + unregister_udfs(_registered) if _corpus: for score in _scores: - if "bleu" in score and not _bleu_registered: + if "bleu" in score: bleu = evaluate.load("bleu") - _bleu_registered = True + _registered.add(f"{response_column}.bleu_score") @register_dataset_udf( [response_column], @@ -42,9 +44,9 @@ def bleu_score(text): ) return result - if "rouge" in score and not _rouge_registered: + if "rouge" in score: rouge = evaluate.load("rouge") - _rouge_registered = True + _registered.add(f"{response_column}.rouge_score") @register_dataset_udf( [response_column], @@ -62,9 +64,9 @@ def rouge_score(text): ) return result - if "meteor" in score and not _meteor_registered: + if "meteor" in score: meteor = evaluate.load("meteor") - _meteor_registered = True + _registered.add(f"{response_column}.meteor_score") @register_dataset_udf( [response_column], @@ -87,11 +89,14 @@ def meteor_score(text): def init( + language: Optional[str] = None, corpus: Optional[str] = None, scores: Set[str] = set(), rouge_type: str = "", config: Optional[LangKitConfig] = None, ): + global _initialized + _initialized = True config = config or deepcopy(lang_config) global _corpus global _scores @@ -101,6 +106,3 @@ def init( _rouge_type = rouge_type or config.rouge_type _register_score_udfs() - - -init() diff --git a/langkit/pattern_loader.py b/langkit/pattern_loader.py index c28e6a29..0b8eba0e 100644 --- a/langkit/pattern_loader.py +++ b/langkit/pattern_loader.py @@ -1,6 +1,5 @@ import json import re -from copy import deepcopy from logging import getLogger from typing import Optional @@ -11,12 +10,14 @@ class PatternLoader: - def __init__(self, config: Optional[LangKitConfig] = None): - self.config: LangKitConfig = config or deepcopy(lang_config) + def __init__(self, json_path: Optional[str] = None): + self.json_path = json_path self.regex_groups = self.load_patterns() def load_patterns(self): - json_path = self.config.pattern_file_path + json_path = self.json_path + if json_path is None: + return None try: skip = False with open(json_path, "r") as myfile: @@ -37,12 +38,10 @@ def load_patterns(self): except json.decoder.JSONDecodeError as json_error: skip = True diagnostic_logger.warning(f"Could not parse {json_path}: {json_error}") - if not skip: - return regex_groups - return None + return regex_groups if not skip else None - def set_config(self, config: LangKitConfig): - self.config = config + def set_config(self, json_path: Optional[str] = None): + self.json_path = json_path def update_patterns(self): self.regex_groups = self.load_patterns() diff --git a/langkit/regexes.py b/langkit/regexes.py index 8f9f6e18..a666a3c4 100644 --- a/langkit/regexes.py +++ b/langkit/regexes.py @@ -6,99 +6,86 @@ from langkit import LangKitConfig, lang_config, prompt_column, response_column from whylogs.core.metrics.metrics import FrequentItemsMetric from whylogs.core.resolvers import MetricSpec -from typing import Dict, List, Optional +from typing import Optional, Set +from langkit.whylogs.unreg import unregister_udfs diagnostic_logger = getLogger(__name__) pattern_loader = PatternLoader() +response_pattern_loader = PatternLoader() +_initialized = False -def has_patterns(text): - regex_groups = pattern_loader.get_regex_groups() + +def has_patterns(text, regex_groups): + if not _initialized: + init() if regex_groups: - matched = None for group in regex_groups: for expression in group["expressions"]: if expression.search(text): - matched = matched or group["name"] - break - if matched is not None: - break - - return matched + return group["name"] + return None -def _wrapper(column): +def _wrapper(column, groups): def wrappee(text): - return [has_patterns(input) for input in text[column]] + return [has_patterns(input, groups) for input in text[column]] return wrappee -_registered: List[str] = [] - - -def _unregister_metric_udf(old_name: str, namespace: Optional[str] = ""): - from whylogs.experimental.core.udf_schema import _multicolumn_udfs - - if _multicolumn_udfs is None or namespace not in _multicolumn_udfs: - return - - _multicolumn_udfs[namespace] = [ - udf - for udf in _multicolumn_udfs[namespace] - if list(udf.udfs.keys())[0] != old_name - ] +_registered: Set[str] = set() def _register_udfs(config: Optional[LangKitConfig] = None): - from whylogs.experimental.core.udf_schema import _resolver_specs - + global _initialized + _initialized = True global _registered - if _registered and config is None: - return + unregister_udfs(_registered) if config is None: config = lang_config default_metric_name = "has_patterns" pattern_metric_name = config.metric_name_map.get( default_metric_name, default_metric_name ) - - for old in _registered: - _unregister_metric_udf(old_name=old) - if ( - _resolver_specs is not None - and isinstance(_resolver_specs, Dict) - and isinstance(_resolver_specs[""], List) - ): - _resolver_specs[""] = [ - spec for spec in _resolver_specs[""] if spec.column_name != old - ] - _registered = [] - if pattern_loader.get_regex_groups() is not None: - for column in [prompt_column, response_column]: - udf_name = f"{column}.{pattern_metric_name}" - register_dataset_udf( - [column], - udf_name=udf_name, - metrics=[MetricSpec(FrequentItemsMetric)], - )(_wrapper(column)) - _registered.append(udf_name) + column = prompt_column + udf_name = f"{column}.{pattern_metric_name}" + register_dataset_udf( + [column], + udf_name=udf_name, + metrics=[MetricSpec(FrequentItemsMetric)], + )(_wrapper(column, pattern_loader.get_regex_groups())) + _registered.add(udf_name) + + if response_pattern_loader.get_regex_groups() is not None: + column = response_column + udf_name = f"{column}.{pattern_metric_name}" + register_dataset_udf( + [column], + udf_name=udf_name, + metrics=[MetricSpec(FrequentItemsMetric)], + )(_wrapper(column, response_pattern_loader.get_regex_groups())) + _registered.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, ): + global _initialized + _initialized = True config = deepcopy(config or lang_config) if pattern_file_path: config.pattern_file_path = pattern_file_path + if response_pattern_file_path: + config.response_pattern_file_path = response_pattern_file_path - global pattern_loader - pattern_loader = PatternLoader(config) - pattern_loader.update_patterns() + global pattern_loader, response_pattern_loader + pattern_loader = PatternLoader(config.pattern_file_path) + response_pattern_loader = PatternLoader(config.response_pattern_file_path) _register_udfs(config) - - -init() diff --git a/langkit/response_hallucination.py b/langkit/response_hallucination.py index 6c2b85d6..367827a2 100644 --- a/langkit/response_hallucination.py +++ b/langkit/response_hallucination.py @@ -2,18 +2,16 @@ from logging import getLogger from typing import List, Optional from whylogs.experimental.core.udf_schema import register_dataset_udf -from langkit import lang_config, prompt_column, response_column +from langkit import LangKitConfig, lang_config, prompt_column, response_column from nltk.tokenize import sent_tokenize from langkit.openai.openai import LLMInvocationParams, Conversation, ChatLog from langkit.transformer import Encoder from sentence_transformers import util -_prompt = prompt_column -_response = response_column diagnostic_logger = getLogger(__name__) -embeddings_encoder = Encoder(lang_config.transformer_name, custom_encoder=None) +embeddings_encoder = None @dataclass @@ -237,24 +235,9 @@ def consistency_check( return consistency_result -checker: Optional[ConsistencyChecker] = None - - -def init(llm: LLMInvocationParams, num_samples=1): - global checker - import nltk - - nltk.download("punkt") - diagnostic_logger.info( - "Info: the response_hallucination metric module performs additionall LLM calls to check the consistency of the response." - ) - checker = ConsistencyChecker(llm, num_samples, embeddings_encoder) - - -@register_dataset_udf([_prompt, _response], f"{_response}.hallucination") def response_hallucination(text): series_result = [] - for prompt, response in zip(text[_prompt], text[_response]): + for prompt, response in zip(text[prompt_column], text[response_column]): result: ConsistencyResult = checker.consistency_check(prompt, response) series_result.append(result.final_score) return series_result @@ -265,3 +248,27 @@ def consistency_check(prompt: str, response: Optional[str] = None): return checker.consistency_check(prompt, response).to_summary_dict() else: raise Exception("You need to call init() before using this function") + + +checker: Optional[ConsistencyChecker] = None + + +def init( + language: Optional[str] = None, + config: Optional[LangKitConfig] = None, + llm: LLMInvocationParams = LLMInvocationParams(), + num_samples=1, +): + config = config or lang_config + global checker, embeddings_encoder + import nltk + + nltk.download("punkt") + diagnostic_logger.info( + "Info: the response_hallucination metric module performs additionall LLM calls to check the consistency of the response." + ) + embeddings_encoder = Encoder(config.response_transformer_name, custom_encoder=None) + checker = ConsistencyChecker(llm, num_samples, embeddings_encoder) + register_dataset_udf( + [prompt_column, response_column], f"{response_column}.hallucination" + )(response_hallucination) diff --git a/langkit/sentiment.py b/langkit/sentiment.py index a92894f2..d2a65542 100644 --- a/langkit/sentiment.py +++ b/langkit/sentiment.py @@ -1,46 +1,167 @@ from copy import deepcopy -from typing import Optional +from typing import Optional, Set from whylogs.experimental.core.udf_schema import register_dataset_udf from langkit import LangKitConfig, lang_config, prompt_column, response_column +from langkit.whylogs.unreg import unregister_udfs -_prompt = prompt_column -_response = response_column +_registered: Set[str] = set() + + +_nltk_downloaded = None +_response_nltk_downloaded = None _sentiment_analyzer = None -_nltk_downloaded = False +_response_sentiment_analyzer = None + +_pipeline = None +_response_pipeline = None +_initialized = False -def sentiment_nltk(text: str) -> float: - if _sentiment_analyzer is None: + +def sentiment_nltk(text: str, sentiment_analyzer=None) -> float: + if not _initialized: + init() + sentiment_analyzer = sentiment_analyzer or _sentiment_analyzer + if sentiment_analyzer is None: raise ValueError( "sentiment metrics must initialize sentiment analyzer before evaluation!" ) - return _sentiment_analyzer.polarity_scores(text)["compound"] + return sentiment_analyzer.polarity_scores(text)["compound"] + +_supported_languages = { + "ar", + "de", + "en", + "es", + "fr", + "hi", + "id", + "it", + "ja", + "ms", + "pt", + "zh", +} -@register_dataset_udf([_prompt], udf_name=f"{_prompt}.sentiment_nltk") -def prompt_sentiment(text): - return [sentiment_nltk(t) for t in text[_prompt]] +def sentiment_multilingual(text: str, pipeline=None) -> float: + if not _initialized: + init() + pipeline = pipeline or _pipeline + if pipeline is None: + raise ValueError("sentiment score must initialize the pipeline first") -@register_dataset_udf([_response], udf_name=f"{_response}.sentiment_nltk") -def response_sentiment(text): - return [sentiment_nltk(t) for t in text[_response]] + result = pipeline(text, return_all_scores=True) + positive = [res["score"] for res in result[0] if res["label"] == "positive"][0] + negative = [res["score"] for res in result[0] if res["label"] == "negative"][0] + return positive - negative -def init(lexicon: Optional[str] = None, config: Optional[LangKitConfig] = None): +def _sentiment_wrapper(sentiment_fn, argument, column): + def _wrappee(text): + return [sentiment_fn(t, argument) for t in text[column]] + + return _wrappee + + +def configure_nltk(config, lexicon, response_lexicon): import nltk from nltk.sentiment import SentimentIntensityAnalyzer - config = config or deepcopy(lang_config) lexicon = lexicon or config.sentiment_lexicon - global _sentiment_analyzer, _nltk_downloaded - if not _nltk_downloaded: + global _nltk_downloaded, _sentiment_analyzer + if _nltk_downloaded != lexicon: + nltk.download(lexicon) + _nltk_downloaded = lexicon + _sentiment_analyzer = ( + SentimentIntensityAnalyzer() + ) # TODO: probably need to pass an argument + else: + _sentiment_analyzer = None + + lexicon = response_lexicon or config.response_sentiment_lexicon + global _response_nltk_downloaded, _response_sentiment_analyzer + if _response_nltk_downloaded != lexicon: nltk.download(lexicon) - _nltk_downloaded = True + _response_nltk_downloaded = lexicon + _response_sentiment_analyzer = ( + SentimentIntensityAnalyzer() + ) # TODO: needs argument + else: + _response_sentiment_analyzer = None + + +def configure_hugging_face(config, sentiment_model_path, response_sentiment_model_path): + from transformers import pipeline + + global _pipeline, _response_pipeline + model_path = sentiment_model_path or config.sentiment_model_path + if model_path: + _pipeline = pipeline(model=model_path, top_k=None) + else: + _pipeline = None + + model_path = response_sentiment_model_path or config.response_sentiment_model_path + if model_path: + _response_pipeline = pipeline(model=model_path, top_k=None) + else: + _response_pipeline = None + + +def init( + language: Optional[str] = None, + lexicon: Optional[str] = None, + config: Optional[LangKitConfig] = None, + response_lexicon: Optional[str] = None, + sentiment_model_path: Optional[str] = None, + response_sentiment_model_path: Optional[str] = None, +): + global _initialized + _initialized = True + + global _registered + unregister_udfs(_registered) + + config = config or deepcopy(lang_config) + prompt_languages = {language} if language is not None else config.prompt_languages + response_languages = ( + {language} if language is not None else config.response_languages + ) - _sentiment_analyzer = SentimentIntensityAnalyzer() + configure_nltk(config, lexicon, response_lexicon) + configure_hugging_face(config, sentiment_model_path, response_sentiment_model_path) + if prompt_languages is not None and len(prompt_languages) > 0: + if prompt_languages.issubset({"", "en"}) and _sentiment_analyzer: + register_dataset_udf( + [prompt_column], udf_name=f"{prompt_column}.sentiment_nltk" + )(_sentiment_wrapper(sentiment_nltk, _sentiment_analyzer, prompt_column)) + _registered.add(f"{prompt_column}.sentiment_nltk") + elif prompt_languages.issubset(_supported_languages) and _pipeline: + register_dataset_udf( + [prompt_column], udf_name=f"{prompt_column}.sentiment_multi" + )(_sentiment_wrapper(sentiment_multilingual, _pipeline, prompt_column)) + _registered.add(f"{prompt_column}.sentiment_multi") -init() + if response_languages is not None and len(response_languages) > 0: + if response_languages.issubset({"", "en"}) and _response_sentiment_analyzer: + register_dataset_udf( + [response_column], udf_name=f"{response_column}.sentiment_nltk" + )( + _sentiment_wrapper( + sentiment_nltk, _response_sentiment_analyzer, response_column + ) + ) + _registered.add(f"{response_column}.sentiment_nltk") + elif response_languages.issubset(_supported_languages) and _response_pipeline: + register_dataset_udf( + [prompt_column], udf_name=f"{response_column}.sentiment_multi" + )( + _sentiment_wrapper( + sentiment_multilingual, _response_pipeline, prompt_column + ) + ) + _registered.add(f"{response_column}.sentiment_multi") diff --git a/langkit/tests/test_injections.py b/langkit/tests/test_injections.py index 633c4937..e5a812be 100644 --- a/langkit/tests/test_injections.py +++ b/langkit/tests/test_injections.py @@ -42,6 +42,7 @@ def test_injections(texts): from langkit import injections # noqa from whylogs.experimental.core.udf_schema import udf_schema + injections.init() text_schema = udf_schema() for text in texts: profile = why.log( diff --git a/langkit/tests/test_textstat.py b/langkit/tests/test_textstat.py index 351a999c..9e609d29 100644 --- a/langkit/tests/test_textstat.py +++ b/langkit/tests/test_textstat.py @@ -24,15 +24,15 @@ def test_textstat(): } ) udf_list = _unpack(ts._udfs_to_register) + [ - ("aggregate_reading_level", ""), + ("aggregate_reading_level", "en"), ] - schema_names = set([s for _, s in udf_list]) - for schema_name in schema_names: - schema = udf_schema(schema_name=schema_name) - view = why.log(df, schema=schema).view() - for stat, stat_schema in udf_list: + languages = set([s for _, s in udf_list]) + for language in languages: + ts.init(language=language) + view = why.log(df, schema=udf_schema()).view() + for stat, lang in udf_list: for column in ["prompt", "response"]: - if stat_schema in {"", schema_name}: + if lang == language: dist = ( view.get_column(f"{column}.{stat}") .get_metric("distribution") diff --git a/langkit/tests/test_themes.py b/langkit/tests/test_themes.py index be038549..25b50965 100644 --- a/langkit/tests/test_themes.py +++ b/langkit/tests/test_themes.py @@ -42,7 +42,7 @@ def test_theme_custom(interactions): def embed(texts: List[str]): return [[0.2, 0.2] for _ in texts] - themes.init(custom_encoder=embed) + themes.init(custom_encoder=embed, response_custom_encoder=embed) schema = udf_schema() for i, interaction in enumerate(interactions): result = why.log(interaction, schema=schema) @@ -129,7 +129,10 @@ def test_themes_with_json_string(): } # if we don't reset udfs, jailbreak_similarity will be an empty metric _reset_udfs() - themes.init(theme_json=json.dumps(refusals_json)) + themes.init( + theme_json=json.dumps(refusals_json), + response_theme_json=json.dumps(refusals_json), + ) schema = udf_schema() prof = why.log({"prompt": "hello"}, schema=schema).view() @@ -150,7 +153,18 @@ def test_themes_with_json_string(): @pytest.mark.load def test_themes_standalone(): - from langkit.themes import group_similarity + from langkit.themes import ( + group_similarity, + init, + _transformer_model, + _embeddings_map, + ) - score = group_similarity("Sorry, but I can't assist with that", "refusal") + init() + score = group_similarity( + "Sorry, but I can't assist with that", + "refusal", + _transformer_model, + _embeddings_map, + ) assert score == pytest.approx(1.0) diff --git a/langkit/tests/test_toxicity.py b/langkit/tests/test_toxicity.py index dc74ce99..57eb3203 100644 --- a/langkit/tests/test_toxicity.py +++ b/langkit/tests/test_toxicity.py @@ -11,6 +11,7 @@ def test_toxicity(): from langkit import toxicity # noqa from whylogs.experimental.core.udf_schema import udf_schema + toxicity.init() text_schema = udf_schema() profile = why.log( {"prompt": "I like you. I love you."}, schema=text_schema @@ -29,6 +30,7 @@ def test_toxicity_long_response(long_response): from langkit import toxicity # noqa from whylogs.experimental.core.udf_schema import udf_schema + toxicity.init() text_schema = udf_schema() profile = why.log(long_response, schema=text_schema).profile() assert profile @@ -39,6 +41,7 @@ def test_empty_toxicity(): from langkit import toxicity # noqa from whylogs.experimental.core.udf_schema import udf_schema + toxicity.init() text_schema = udf_schema() test_prompt = "hi." test_response = "" diff --git a/langkit/tests/test_translation.py b/langkit/tests/test_translation.py new file mode 100644 index 00000000..caa1e24b --- /dev/null +++ b/langkit/tests/test_translation.py @@ -0,0 +1,67 @@ +from typing import Dict, List, Union +from langkit.translator import Translator, translated, translated_udf +import whylogs as why +from whylogs.core.metrics import FrequentItemsMetric +from whylogs.core.resolvers import MetricSpec +from whylogs.core.stubs import pd as pd +from whylogs.experimental.core.udf_schema import register_dataset_udf, udf_schema + + +class FakeTranslator(Translator): + def translate(self, text: str) -> str: + return text[::-1] + + +class NonTranslator(Translator): + def translate(self, text: str) -> str: + return text + + +@translated(FakeTranslator()) +def translated_prompt(x): + return x + + +def wrapper(input): + return [translated_prompt(x) for x in input["prompt"]] + + +def test_translated_decorator(): + register_dataset_udf( + ["prompt"], + udf_name="prompt.translated", + metrics=[MetricSpec(FrequentItemsMetric)], + schema_name="translated", + )(wrapper) + untranslated = {"prompt": "This is a prompt", "response": "This is a response"} + schema = udf_schema(schema_name="translated", include_default_schema=False) + result = why.log(untranslated, schema=schema) + assert ( + result.view() + .get_column("prompt.translated") + .get_metric("frequent_items") + .strings[0] + .value + == untranslated["prompt"][::-1] + ) + + +@translated_udf({"prompt": NonTranslator(), "response": FakeTranslator()}) +def concat(input: Union[Dict[str, List], pd.DataFrame]) -> Union[List, pd.Series]: + return [x + y for x, y in zip(input["prompt"], input["response"])] + + +def test_translated_udf(): + register_dataset_udf( + ["prompt", "response"], + udf_name="bob", + metrics=[MetricSpec(FrequentItemsMetric)], + schema_name="translated", + )(concat) + untranslated = {"prompt": "knock knock", "response": "who's there?"} + schema = udf_schema(schema_name="translated", include_default_schema=False) + result = why.log(untranslated, schema=schema) + assert ( + result.view().get_column("bob").get_metric("frequent_items").strings[0].value + == untranslated["prompt"] + untranslated["response"][::-1] + ) diff --git a/langkit/textstat.py b/langkit/textstat.py index deca402a..a2ee0c2d 100644 --- a/langkit/textstat.py +++ b/langkit/textstat.py @@ -1,8 +1,10 @@ from logging import getLogger -from typing import Callable, Dict, List, Optional, Tuple, Union +from typing import Callable, Dict, List, Optional, Set, Tuple, Union +import textstat from whylogs.core.stubs import pd from whylogs.experimental.core.udf_schema import register_dataset_udf -from langkit import LangKitConfig, prompt_column, response_column +from langkit import LangKitConfig, lang_config, prompt_column, response_column +from langkit.whylogs.unreg import unregister_udfs diagnostic_logger = getLogger(__name__) @@ -10,17 +12,19 @@ # score metrics -# (stat name, schema name[, udf name]) +# TODO: should probably s/""/"en"/ + +# (stat name, language[, udf name]) _udfs_to_register: List[Union[Tuple[str, str], Tuple[str, str, str]]] = [ ("flesch_kincaid_grade", "text_standard_component"), - ("flesch_reading_ease", ""), + ("flesch_reading_ease", "en"), ("smog_index", "text_standard_component"), ("coleman_liau_index", "text_standard_component"), - ("automated_readability_index", ""), + ("automated_readability_index", "en"), ("dale_chall_readability_score", "text_standard_component"), ("linsear_write_formula", "text_standard_component"), ("gunning_fog", "text_standard_component"), - ("text_standard", "", "aggregate_reading_level"), + ("text_standard", "en", "aggregate_reading_level"), ("fernandez_huerta", "es"), ("szigriszt_pazos", "es"), ("gutierrez_polini", "es"), @@ -28,14 +32,14 @@ ("gulpease_index", "it"), ("osman", "ar"), # count metrics - ("syllable_count", ""), - ("lexicon_count", ""), - ("sentence_count", ""), - ("char_count", "", "character_count"), - ("letter_count", ""), - ("polysyllabcount", "", "polysyllable_count"), - ("monosyllabcount", "", "monosyllable_count"), - ("difficult_words", ""), + ("syllable_count", "en"), + ("lexicon_count", "en"), + ("sentence_count", "en"), + ("char_count", "en", "character_count"), + ("letter_count", "en"), + ("polysyllabcount", "en", "polysyllable_count"), + ("monosyllabcount", "en", "monosyllable_count"), + ("difficult_words", "en"), ] @@ -65,31 +69,43 @@ def wrappee(text: Union[pd.DataFrame, Dict[str, List]]) -> Union[pd.Series, List return wrappee -def init(config: Optional[LangKitConfig] = None): - pass - - -init() - - def _unpack(t: Union[Tuple[str, str], Tuple[str, str, str]]) -> Tuple[str, str, str]: return t if len(t) == 3 else (t[0], t[1], t[0]) # type: ignore -_registered = False +_registered: Set[str] = set() + +def init(language: Optional[str] = None, config: Optional[LangKitConfig] = None): + config = config or lang_config + prompt_languages = ( + {language} if language is not None else config.prompt_languages + ) or set() + response_languages = ( + {language} if language is not None else config.response_languages + ) or set() + global _registered + unregister_udfs(_registered) -if not _registered: - _registered = True for t in _udfs_to_register: stat_name, schema_name, udf = _unpack(t) for column in [prompt_column, response_column]: - register_dataset_udf( - [column], udf_name=f"{column}.{udf}", schema_name=schema_name - )(wrapper(stat_name, column)) + if schema_name in ( + prompt_languages if column == prompt_column else response_languages + ): + udf_name = f"{column}.{udf}" + register_dataset_udf( + [column], + udf_name=udf_name, + # schema_name=schema_name, # TODO: probably should be default schema + )(wrapper(stat_name, column)) + _registered.add(udf_name) for column in [prompt_column, response_column]: - register_dataset_udf([column], udf_name=f"{column}.aggregate_reading_level")( - aggregate_wrapper(column) - ) + if "en" in ( + prompt_languages if column == prompt_column else response_languages + ): + udf_name = f"{column}.aggregate_reading_level" + register_dataset_udf([column], udf_name=udf_name)(aggregate_wrapper(column)) + _registered.add(udf_name) diagnostic_logger.info("Initialized textstat metrics.") diff --git a/langkit/themes.py b/langkit/themes.py index 4639081f..80ea7e22 100644 --- a/langkit/themes.py +++ b/langkit/themes.py @@ -1,7 +1,7 @@ import json from copy import deepcopy from logging import getLogger -from typing import Callable, Optional, Dict, List +from typing import Callable, Optional, Dict, List, Set from sentence_transformers import util from torch import Tensor @@ -10,67 +10,95 @@ from langkit.transformer import Encoder from langkit import LangKitConfig, lang_config, prompt_column, response_column +from langkit.whylogs.unreg import unregister_udfs + diagnostic_logger = getLogger(__name__) _transformer_model = None _theme_groups = None -_prompt = prompt_column -_response = response_column - _embeddings_map: Dict[str, List] = {} -def create_similarity_function(group: str, column: str): +_response_transformer_model = None +_response_theme_groups = None +_response_embeddings_map: Dict[str, List] = {} + + +def create_similarity_function( + group: str, column: str, transformer_model, embeddings_map: Dict[str, List] +): def similarity_by_group(text): result = [] for input in text[column]: - score = group_similarity(input, group) + score = group_similarity(input, group, transformer_model, embeddings_map) result.append(score) return result return similarity_by_group -def group_similarity(text: str, group): +def group_similarity( + text: str, group, transformer_model, embeddings_map: Dict[str, List] +): similarities: List[float] = [] - if _transformer_model is None: + if transformer_model is None: raise ValueError("Must initialize a transformer before calling encode!") - text_embedding = _transformer_model.encode(text) - for embedding in _embeddings_map.get(group, []): + text_embedding = transformer_model.encode(text) + for embedding in embeddings_map.get(group, []): similarity = get_embeddings_similarity(text_embedding, embedding) similarities.append(similarity) return max(similarities) if similarities else None -def _map_embeddings(): - global _embeddings_map - for group in _theme_groups: - _embeddings_map[group] = [ - _transformer_model.encode(s) for s in _theme_groups.get(group, []) +def _map_embeddings(embeddings_map, theme_groups, transformer_model): + for group in theme_groups: + embeddings_map[group] = [ + transformer_model.encode(s) for s in theme_groups.get(group, []) ] -_registered = set() +_registered: Set[str] = set() def _register_theme_udfs(): global _registered - _map_embeddings() - - for group in _theme_groups: - for column in [_prompt, _response]: - if group == "jailbreak" and column == _response: + unregister_udfs(_registered) + if _transformer_model is not None: + _map_embeddings(_embeddings_map, _theme_groups, _transformer_model) + for group in _theme_groups: + column = prompt_column + if group == "refusal": continue - if group == "refusal" and column == _prompt: + udf_name = f"{column}.{group}_similarity" + _registered.add(udf_name) + register_dataset_udf([column], udf_name=udf_name)( + create_similarity_function( + group, column, _transformer_model, _embeddings_map + ) + ) + + if _response_transformer_model is not None: + _map_embeddings( + _response_embeddings_map, + _response_theme_groups, + _response_transformer_model, + ) + for group in _response_theme_groups: + column = response_column + if group == "jailbreak": continue udf_name = f"{column}.{group}_similarity" - if udf_name not in _registered: - _registered.add(udf_name) - register_dataset_udf([column], udf_name=udf_name)( - create_similarity_function(group, column) + _registered.add(udf_name) + register_dataset_udf([column], udf_name=udf_name)( + create_similarity_function( + group, + column, + _response_transformer_model, + _response_embeddings_map, ) + ) def load_themes(json_path: str, encoding="utf-8"): @@ -90,34 +118,71 @@ def load_themes(json_path: str, encoding="utf-8"): def init( + language: Optional[str] = None, transformer_name: Optional[str] = None, custom_encoder: Optional[Callable] = None, theme_file_path: Optional[str] = None, theme_json: Optional[str] = None, config: Optional[LangKitConfig] = None, + response_transformer_name: Optional[str] = None, + response_custom_encoder: Optional[Callable] = None, + response_theme_file_path: Optional[str] = None, + response_theme_json: Optional[str] = None, ): config = config or deepcopy(lang_config) global _transformer_model global _theme_groups if not transformer_name and not custom_encoder: transformer_name = config.transformer_name - _transformer_model = Encoder(transformer_name, custom_encoder) + if not transformer_name and not custom_encoder: + _transformer_model = None + else: + _transformer_model = Encoder(transformer_name, custom_encoder) + if theme_file_path is not None and theme_json is not None: raise ValueError("Cannot specify both theme_file_path and theme_json") - if theme_file_path is None: - if theme_json: - _theme_groups = json.loads(theme_json) - else: - _theme_groups = load_themes(config.theme_file_path) - else: + + theme_file_path = theme_file_path or config.theme_file_path + if theme_json: + _theme_groups = json.loads(theme_json) + elif theme_file_path: _theme_groups = load_themes(theme_file_path) + else: + _transformer_model = None + + global _response_transformer_model + global _response_theme_groups + if not response_transformer_name and not response_custom_encoder: + response_transformer_name = config.response_transformer_name + if not response_transformer_name and not response_custom_encoder: + _response_transformer_model = None + else: + _response_transformer_model = Encoder( + response_transformer_name, response_custom_encoder + ) + if response_theme_file_path is not None and response_theme_json is not None: + raise ValueError( + "Cannot specify both response_theme_file_path and response_theme_json" + ) + response_theme_file_path = ( + response_theme_file_path or config.response_theme_file_path + ) + if response_theme_json: + _response_theme_groups = json.loads(response_theme_json) + elif response_theme_file_path: + _response_theme_groups = load_themes(response_theme_file_path) + else: + _response_transformer_model = None + _register_theme_udfs() -def get_subject_similarity(text: str, comparison_embedding: Tensor) -> float: - if _transformer_model is None: +def get_subject_similarity( + text: str, comparison_embedding: Tensor, transformer_model +) -> float: + if transformer_model is None: raise ValueError("Must initialize a transformer before calling encode!") - embedding = _transformer_model.encode(text) + embedding = transformer_model.encode(text) similarity = util.pytorch_cos_sim(embedding, comparison_embedding) return similarity.item() @@ -125,10 +190,5 @@ def get_subject_similarity(text: str, comparison_embedding: Tensor) -> float: def get_embeddings_similarity( text_embedding: Tensor, comparison_embedding: Tensor ) -> float: - if _transformer_model is None: - raise ValueError("Must initialize a transformer before calling encode!") similarity = util.pytorch_cos_sim(text_embedding, comparison_embedding) return similarity.item() - - -init() diff --git a/langkit/topics.py b/langkit/topics.py index 7f9bb058..d6044f44 100644 --- a/langkit/topics.py +++ b/langkit/topics.py @@ -1,10 +1,11 @@ from copy import deepcopy from whylogs.experimental.core.udf_schema import register_dataset_udf -from typing import Callable, List, Optional +from typing import Callable, List, Optional, Set from transformers import ( pipeline, ) from langkit import LangKitConfig, lang_config, prompt_column, response_column +from langkit.whylogs.unreg import unregister_udfs import os import torch @@ -16,35 +17,76 @@ _topics: List[str] = lang_config.topics +_model_path: Optional[str] = None +_classifier = None -_model_path: str = lang_config.topic_model_path -_classifier = pipeline(lang_config.topic_classifier, model=_model_path, device=_device) +_response_topics: List[str] = lang_config.response_topics +_response_model_path: Optional[str] = None +_response_classifier = None +_initialized = False -def closest_topic(text): - return _classifier(text, _topics, multi_label=False)["labels"][0] +def closest_topic(text, classifier=None, topics=None): + if not _initialized: + init() + classifier = classifier or _classifier + topics = topics or _topics + if classifier is None: + raise ValueError("Topics - classifier model not initialized") + return classifier(text, topics, multi_label=False)["labels"][0] -def _wrapper(column: str) -> Callable: - return lambda text: [closest_topic(t) for t in text[column]] + +def _wrapper(column: str, classifier, topics) -> Callable: + return lambda text: [closest_topic(t, classifier, topics) for t in text[column]] + + +_registered: Set[str] = set() def init( + language: Optional[str] = None, topics: Optional[List[str]] = None, model_path: Optional[str] = None, topic_classifier: Optional[str] = None, config: Optional[LangKitConfig] = None, + response_topics: Optional[List[str]] = None, + response_model_path: Optional[str] = None, + response_topic_classifier: Optional[str] = None, ): + global _initialized + _initialized = True + global _registered + unregister_udfs(_registered) config = config or deepcopy(lang_config) global _topics, _classifier _topics = topics or config.topics topic_classifier = topic_classifier or lang_config.topic_classifier model_path = model_path or config.topic_model_path - _classifier = pipeline(topic_classifier, model=model_path, device=_device) - for column in [prompt_column, response_column]: - register_dataset_udf([column], udf_name=f"{column}.closest_topic")( - _wrapper(column) - ) + if not (model_path and topic_classifier): + _classifier = None + else: + _classifier = pipeline(topic_classifier, model=model_path, device=_device) + + global _response_topics, _response_classifier + _response_topics = response_topics or config.response_topics + topic_classifier = ( + response_topic_classifier or lang_config.response_topic_classifier + ) + model_path = response_model_path or config.response_topic_model_path + if not (model_path and topic_classifier): + _response_classifier = None + else: + _response_classifier = pipeline(topic_classifier, model=model_path, device=_device) + if _classifier is not None: + register_dataset_udf( + [prompt_column], udf_name=f"{prompt_column}.closest_topic" + )(_wrapper(prompt_column, _classifier, _topics)) + _registered.add(f"{prompt_column}.closest_topic") -init() + if _response_classifier is not None: + register_dataset_udf( + [response_column], udf_name=f"{response_column}.closest_topic" + )(_wrapper(response_column, _response_classifier, _response_topics)) + _registered.add(f"{response_column}.closest_topic") diff --git a/langkit/toxicity.py b/langkit/toxicity.py index 8e1f1979..22034ef6 100644 --- a/langkit/toxicity.py +++ b/langkit/toxicity.py @@ -1,8 +1,10 @@ from copy import deepcopy -from typing import Optional +from typing import Optional, Set from whylogs.experimental.core.udf_schema import register_dataset_udf from langkit import LangKitConfig, lang_config, prompt_column, response_column +from langkit.translator import Translator, translated_udf +from langkit.whylogs.unreg import unregister_udfs import os import torch @@ -12,49 +14,97 @@ ) _device = 0 if _USE_CUDA else -1 -_prompt = prompt_column -_response = response_column _toxicity_tokenizer = None _toxicity_pipeline = None +_response_toxicity_tokenizer = None +_response_toxicity_pipeline = None -def toxicity(text: str) -> float: - if _toxicity_pipeline is None or _toxicity_tokenizer is None: +_initialized = False + +PROMPT_TRANSLATOR: Optional[Translator] = None +RESPONSE_TRANSLATOR: Optional[Translator] = None +TRANSLATOR: Optional[Translator] = None + + +def toxicity(text: str, pipeline, tokenizer) -> float: + if not _initialized: + init() + pipeline = pipeline or _toxicity_pipeline + tokenizer = tokenizer or _toxicity_tokenizer + if pipeline is None or tokenizer is None: raise ValueError("toxicity score must initialize the pipeline first") - result = _toxicity_pipeline( - text, truncation=True, max_length=_toxicity_tokenizer.model_max_length - ) + result = pipeline(text, truncation=True, max_length=tokenizer.model_max_length) return ( result[0]["score"] if result[0]["label"] == "toxic" else 1 - result[0]["score"] ) -@register_dataset_udf([_prompt], f"{_prompt}.toxicity") -def prompt_toxicity(text): - return [toxicity(t) for t in text[_prompt]] +def _toxicity_wrapper(column, pipeline, tokenizer): + return lambda text: [toxicity(t, pipeline, tokenizer) for t in text[column]] -@register_dataset_udf([_response], f"{_response}.toxicity") -def response_toxicity(text): - return [toxicity(t) for t in text[_response]] +_registered: Set[str] = set() -def init(model_path: Optional[str] = None, config: Optional[LangKitConfig] = None): +def init( + language: Optional[str] = None, + model_path: Optional[str] = None, + config: Optional[LangKitConfig] = None, + response_model_path: Optional[str] = None, +): + global _initialized + _initialized = True + global _registered + unregister_udfs(_registered) from transformers import ( AutoModelForSequenceClassification, AutoTokenizer, TextClassificationPipeline, ) + translators = { + prompt_column: TRANSLATOR or PROMPT_TRANSLATOR, + response_column: TRANSLATOR or RESPONSE_TRANSLATOR, + } config = config or deepcopy(lang_config) model_path = model_path or config.toxicity_model_path global _toxicity_tokenizer, _toxicity_pipeline - _toxicity_tokenizer = AutoTokenizer.from_pretrained(model_path) - model = AutoModelForSequenceClassification.from_pretrained(model_path) - _toxicity_pipeline = TextClassificationPipeline( - model=model, tokenizer=_toxicity_tokenizer, device=_device - ) - - -init() + if model_path is None: + _toxicity_pipeline = None + else: + _toxicity_tokenizer = AutoTokenizer.from_pretrained(model_path) + model = AutoModelForSequenceClassification.from_pretrained(model_path) + _toxicity_pipeline = TextClassificationPipeline( + model=model, tokenizer=_toxicity_tokenizer, device=_device + ) + register_dataset_udf([prompt_column], f"{prompt_column}.toxicity")( + translated_udf(translators)( + _toxicity_wrapper( + prompt_column, _toxicity_pipeline, _toxicity_tokenizer + ) + ) + ) + _registered.add(f"{prompt_column}.toxicity") + + model_path = response_model_path or config.response_toxicity_model_path + global _response_toxicity_tokenizer, _response_toxicity_pipeline + if model_path is None: + _response_toxicity_pipeline = None + else: + _response_toxicity_tokenizer = AutoTokenizer.from_pretrained(model_path) + model = AutoModelForSequenceClassification.from_pretrained(model_path) + _response_toxicity_pipeline = TextClassificationPipeline( + model=model, tokenizer=_response_toxicity_tokenizer, device=_device + ) + register_dataset_udf([response_column], f"{response_column}.toxicity")( + translated_udf(translators)( + _toxicity_wrapper( + response_column, + _response_toxicity_pipeline, + _response_toxicity_tokenizer, + ) + ) + ) + _registered.add(f"{response_column}.toxicity") diff --git a/langkit/translator.py b/langkit/translator.py new file mode 100644 index 00000000..0a734118 --- /dev/null +++ b/langkit/translator.py @@ -0,0 +1,48 @@ +from abc import ABC, abstractmethod +from typing import Any, Callable, Dict, List, Optional, Union +from whylogs.core.stubs import pd as pd + + +class Translator(ABC): + @abstractmethod + def translate(self, text: str) -> str: + return text + + +def translated(translator: Optional[Translator] = None): + def decorator(func: Callable[[str], Any]) -> Callable[[str], Any]: + return lambda text: (func(translator.translate(text)) if translator else text) + + return decorator + + +def translated_udf(translators: Optional[Dict[str, Optional[Translator]]]): + def decorator( + func: Callable[[Union[Dict[str, List], pd.DataFrame]], Union[List, pd.Series]] + ) -> Callable[[Union[Dict[str, List], pd.DataFrame]], Union[List, pd.Series]]: + def wrapper( + text: Union[Dict[str, List], pd.DataFrame] + ) -> Union[List, pd.Series]: + if translators is None: + return text + + if isinstance(text, dict): + translated = { + k: [ + (translators[k].translate(t) if translators[k] else t) # type: ignore + for t in v + ] + for k, v in text.items() + } + return func(translated) + + translated = pd.DataFrame() + for k in text.keys(): + translated[k] = ( + text[k].map(translators[k].translate) if translators[k] else text[k] # type: ignore + ) + return func(translated) + + return wrapper + + return decorator diff --git a/langkit/whylogs/unreg.py b/langkit/whylogs/unreg.py new file mode 100644 index 00000000..aa203347 --- /dev/null +++ b/langkit/whylogs/unreg.py @@ -0,0 +1,35 @@ +import whylogs.experimental.core.udf_schema as us +from typing import Optional, Set +import logging + +logger = logging.getLogger(__name__) + + +def unregister_udf( + udf_name: str, namespace: Optional[str] = None, schema_name: str = "" +) -> None: + name = f"{namespace}.{udf_name}" if namespace else udf_name + if schema_name not in us._multicolumn_udfs: + logger.warn( + f"Can't unregister UDF {name} from non-existant schema {schema_name}" + ) + return + + found = False + for spec in us._multicolumn_udfs[schema_name]: + if name in spec.udfs: + found = True + del spec.udfs[name] + if not found: + logger.warn(f"UDF {name} could not be found for unregistering") + us._resolver_specs[schema_name] = list( + filter(lambda x: x.column_name != name, us._resolver_specs[schema_name]) + ) + + +def unregister_udfs( + udfs: Set[str], namespace: Optional[str] = None, schema_name: str = "" +) -> None: + for udf in udfs: + unregister_udf(udf, namespace, schema_name) + udfs = set() diff --git a/pyproject.toml b/pyproject.toml index 15c0da92..cad51d04 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langkit" -version = "0.0.25" +version = "0.0.26-dev1" description = "A collection of text metric udfs for whylogs profiling and monitoring in WhyLabs" authors = ["WhyLabs.ai "] homepage = "https://docs.whylabs.ai/docs/large-language-model-monitoring" From bf996527a8f15bc6dcd91cf1168cb92017e2ffca Mon Sep 17 00:00:00 2001 From: Richard Rogers Date: Wed, 22 Nov 2023 05:52:09 +0000 Subject: [PATCH 2/3] pre-commit cleanup after conflict resolution from rebase --- langkit/injections.py | 1 + langkit/pattern_loader.py | 2 -- langkit/textstat.py | 1 - langkit/topics.py | 4 +++- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/langkit/injections.py b/langkit/injections.py index a38728f8..324569d7 100644 --- a/langkit/injections.py +++ b/langkit/injections.py @@ -19,6 +19,7 @@ _initialized = False + def injection(prompt: Union[Dict[str, List], pd.DataFrame]) -> Union[List, pd.Series]: if not _initialized: init() diff --git a/langkit/pattern_loader.py b/langkit/pattern_loader.py index 0b8eba0e..3071f47f 100644 --- a/langkit/pattern_loader.py +++ b/langkit/pattern_loader.py @@ -3,8 +3,6 @@ from logging import getLogger from typing import Optional -from langkit import LangKitConfig, lang_config - diagnostic_logger = getLogger(__name__) diff --git a/langkit/textstat.py b/langkit/textstat.py index a2ee0c2d..977e8400 100644 --- a/langkit/textstat.py +++ b/langkit/textstat.py @@ -1,6 +1,5 @@ from logging import getLogger from typing import Callable, Dict, List, Optional, Set, Tuple, Union -import textstat from whylogs.core.stubs import pd from whylogs.experimental.core.udf_schema import register_dataset_udf from langkit import LangKitConfig, lang_config, prompt_column, response_column diff --git a/langkit/topics.py b/langkit/topics.py index d6044f44..fa549ec1 100644 --- a/langkit/topics.py +++ b/langkit/topics.py @@ -77,7 +77,9 @@ def init( if not (model_path and topic_classifier): _response_classifier = None else: - _response_classifier = pipeline(topic_classifier, model=model_path, device=_device) + _response_classifier = pipeline( + topic_classifier, model=model_path, device=_device + ) if _classifier is not None: register_dataset_udf( From df0f6ad0743f9b75d91586f43a83d31d4741b2a8 Mon Sep 17 00:00:00 2001 From: Richard Rogers Date: Sat, 2 Dec 2023 01:39:40 -0800 Subject: [PATCH 3/3] specify schema name to register UDFs into --- langkit/all_metrics.py | 40 +++++++++++++++++++++++-------- langkit/count_regexes.py | 24 ++++++++++++------- langkit/injections.py | 15 ++++++++---- langkit/input_output.py | 14 +++++++---- langkit/light_metrics.py | 16 +++++++++---- langkit/llm_metrics.py | 32 ++++++++++++++++++------- langkit/nlp_scores.py | 24 ++++++++++++------- langkit/regexes.py | 21 ++++++++++------ langkit/response_hallucination.py | 18 +++++++++++--- langkit/sentiment.py | 35 ++++++++++++++++++--------- langkit/textstat.py | 29 +++++++++++++++------- langkit/themes.py | 21 +++++++++------- langkit/topics.py | 23 ++++++++++++------ langkit/toxicity.py | 23 ++++++++++++------ langkit/whylogs/unreg.py | 5 ++++ 15 files changed, 242 insertions(+), 98 deletions(-) diff --git a/langkit/all_metrics.py b/langkit/all_metrics.py index 6e280c55..e6626d33 100644 --- a/langkit/all_metrics.py +++ b/langkit/all_metrics.py @@ -16,15 +16,35 @@ def init( - language: Optional[str] = None, config: Optional[LangKitConfig] = None + language: Optional[str] = None, + config: Optional[LangKitConfig] = None, + schema_name: str = "", ) -> DeclarativeSchema: - injections.init(language, config=config or multi_lang_config[language]) - topics.init(language, config=config or multi_lang_config[language]) - regexes.init(language, config=config or multi_lang_config[language]) - sentiment.init(language, config=config or multi_lang_config[language]) - textstat.init(language, config=config or multi_lang_config[language]) - themes.init(language, config=config or multi_lang_config[language]) - toxicity.init(language, config=config or multi_lang_config[language]) - input_output.init(language, config=config or multi_lang_config[language]) - text_schema = attach_schema_metadata(udf_schema(), "all_metrics") + 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 diff --git a/langkit/count_regexes.py b/langkit/count_regexes.py index fa12bf26..21c0d68b 100644 --- a/langkit/count_regexes.py +++ b/langkit/count_regexes.py @@ -1,3 +1,4 @@ +from collections import defaultdict from copy import deepcopy from logging import getLogger @@ -35,12 +36,16 @@ def wrappee(text: Union[pd.DataFrame, Dict[str, List]]) -> Union[pd.Series, List return wrappee -_registered: Set[str] = set() +_registered: Dict[str, Set[str]] = defaultdict( + set +) # _registered[schema_name] -> set of registered UDF names -def _register_udfs(language: str): +def _register_udfs(language: str, schema_name: str): global _registered - unregister_udfs(_registered) + 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: column = prompt_column @@ -49,9 +54,10 @@ def _register_udfs(language: str): register_dataset_udf( [column], udf_name=udf_name, - schema_name=language, + namespace=language, + schema_name=schema_name, )(wrapper(group, column)) - _registered.add(udf_name) + _registered[schema_name].add(udf_name) regex_groups = response_pattern_loader.get_regex_groups() if regex_groups is not None: @@ -61,9 +67,10 @@ def _register_udfs(language: str): register_dataset_udf( [column], udf_name=udf_name, - schema_name=language, + namespace=language, + schema_name=schema_name, )(wrapper(group, column)) - _registered.add(udf_name) + _registered[schema_name].add(udf_name) def init( @@ -71,6 +78,7 @@ def init( pattern_file_path: Optional[str] = None, config: Optional[LangKitConfig] = None, response_pattern_file_path: Optional[str] = None, + schema_name: str = "", ): global _initialized _initialized = True @@ -83,4 +91,4 @@ def init( 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) + _register_udfs(language, schema_name) diff --git a/langkit/injections.py b/langkit/injections.py index 324569d7..9d326180 100644 --- a/langkit/injections.py +++ b/langkit/injections.py @@ -1,3 +1,4 @@ +from collections import defaultdict from copy import deepcopy from typing import Dict, List, Optional, Set, Union from whylogs.core.stubs import pd @@ -48,7 +49,9 @@ def download_embeddings(url): return array -_registered: Set[str] = set() +_registered: Dict[str, Set[str]] = defaultdict( + set +) # _registered[schema_name] -> set of registered UDF names def init( @@ -56,11 +59,13 @@ def init( transformer_name: Optional[str] = None, version: Optional[str] = None, config: Optional[LangKitConfig] = None, + schema_name: str = "", ): global _initialized _initialized = True global _registered - unregister_udfs(_registered) + unregister_udfs(_registered[schema_name], schema_name=schema_name) + _registered[schema_name] = set() config = config or deepcopy(lang_config) global _transformer_model global _index_embeddings @@ -109,5 +114,7 @@ def init( f"Injections - unable to deserialize index to {embeddings_path}. Error: {deserialization_error}" ) if _index_embeddings and _transformer_model: - register_dataset_udf([prompt_column], f"{prompt_column}.injection")(injection) - _registered.add(f"{prompt_column}.injection") + register_dataset_udf( + [prompt_column], f"{prompt_column}.injection", schema_name=schema_name + )(injection) + _registered[schema_name].add(f"{prompt_column}.injection") diff --git a/langkit/input_output.py b/langkit/input_output.py index 8a580263..fc9ac74d 100644 --- a/langkit/input_output.py +++ b/langkit/input_output.py @@ -1,6 +1,7 @@ +from collections import defaultdict from copy import deepcopy from logging import getLogger -from typing import Callable, Optional, Set +from typing import Callable, Dict, Optional, Set from sentence_transformers import util from whylogs.experimental.core.udf_schema import register_dataset_udf @@ -44,7 +45,9 @@ def prompt_response_similarity(text): return series_result -_registered: Set[str] = set() +_registered: Dict[str, Set[str]] = defaultdict( + set +) # _registered[schema_name] -> set of registered UDF names def init( @@ -52,11 +55,13 @@ def init( transformer_name: Optional[str] = None, custom_encoder: Optional[Callable] = None, config: Optional[LangKitConfig] = None, + schema_name: str = "", ): global _initialized _initialized = True global _registered - unregister_udfs(_registered) + unregister_udfs(_registered[schema_name], schema_name=schema_name) + _registered[schema_name] = set() if transformer_name and custom_encoder: raise ValueError( "Only one of transformer_name or encoder can be specified, not both." @@ -81,5 +86,6 @@ def init( register_dataset_udf( [prompt_column, response_column], f"{response_column}.relevance_to_{prompt_column}", + schema_name=schema_name, )(prompt_response_similarity) - _registered.add(f"{response_column}.relevance_to_{prompt_column}") + _registered[schema_name].add(f"{response_column}.relevance_to_{prompt_column}") diff --git a/langkit/light_metrics.py b/langkit/light_metrics.py index c8750165..a2568c26 100644 --- a/langkit/light_metrics.py +++ b/langkit/light_metrics.py @@ -9,10 +9,18 @@ def init( - language: Optional[str] = None, config: Optional[LangKitConfig] = None + language: Optional[str] = None, + config: Optional[LangKitConfig] = None, + schema_name: str = "", ) -> DeclarativeSchema: - regexes.init(language, config=config or multi_lang_config[language]) - textstat.init(language, config=config or multi_lang_config[language]) + regexes.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 + ) - text_schema = attach_schema_metadata(udf_schema(), "light_metrics") + text_schema = attach_schema_metadata( + udf_schema(schema_name=schema_name), "light_metrics" + ) return text_schema diff --git a/langkit/llm_metrics.py b/langkit/llm_metrics.py index 062e79e1..cafe450d 100644 --- a/langkit/llm_metrics.py +++ b/langkit/llm_metrics.py @@ -21,14 +21,30 @@ def init( - language: Optional[str] = None, config: Optional[LangKitConfig] = None + language: Optional[str] = None, + config: Optional[LangKitConfig] = None, + schema_name: str = "", ) -> DeclarativeSchema: - regexes.init(language, config=config or multi_lang_config[language]) - sentiment.init(language, config=config or multi_lang_config[language]) - textstat.init(language, config=config or multi_lang_config[language]) - themes.init(language, config=config or multi_lang_config[language]) - toxicity.init(language, config=config or multi_lang_config[language]) - input_output.init(language, config=config or multi_lang_config[language]) + 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(), "llm_metrics") + text_schema = attach_schema_metadata( + udf_schema(schema_name=schema_name), "llm_metrics" + ) return text_schema diff --git a/langkit/nlp_scores.py b/langkit/nlp_scores.py index e62397c6..b8b873d0 100644 --- a/langkit/nlp_scores.py +++ b/langkit/nlp_scores.py @@ -1,5 +1,6 @@ +from collections import defaultdict from copy import deepcopy -from typing import List, Optional, Set +from typing import Dict, List, Optional, Set from whylogs.experimental.core.udf_schema import register_dataset_udf import evaluate from langkit import LangKitConfig, lang_config, response_column @@ -16,23 +17,27 @@ _initialized = False -_registered: Set[str] = set() +_registered: Dict[str, Set[str]] = defaultdict( + set +) # _registered[schema_name] -> set of registered UDF names -def _register_score_udfs(): +def _register_score_udfs(schema_name: str): if not _initialized: init() global _registered - unregister_udfs(_registered) + unregister_udfs(_registered[schema_name], schema_name=schema_name) + _registered[schema_name] = set() if _corpus: for score in _scores: if "bleu" in score: bleu = evaluate.load("bleu") - _registered.add(f"{response_column}.bleu_score") + _registered[schema_name].add(f"{response_column}.bleu_score") @register_dataset_udf( [response_column], udf_name=f"{response_column}.bleu_score", + schema_name=schema_name, ) def bleu_score(text): result = [] @@ -46,11 +51,12 @@ def bleu_score(text): if "rouge" in score: rouge = evaluate.load("rouge") - _registered.add(f"{response_column}.rouge_score") + _registered[schema_name].add(f"{response_column}.rouge_score") @register_dataset_udf( [response_column], udf_name=f"{response_column}.rouge_score", + schema_name=schema_name, ) def rouge_score(text): result = [] @@ -66,11 +72,12 @@ def rouge_score(text): if "meteor" in score: meteor = evaluate.load("meteor") - _registered.add(f"{response_column}.meteor_score") + _registered[schema_name].add(f"{response_column}.meteor_score") @register_dataset_udf( [response_column], udf_name=f"{response_column}.meteor_score", + schema_name=schema_name, ) def meteor_score(text): result = [] @@ -94,6 +101,7 @@ def init( scores: Set[str] = set(), rouge_type: str = "", config: Optional[LangKitConfig] = None, + schema_name: str = "", ): global _initialized _initialized = True @@ -105,4 +113,4 @@ def init( _scores = list(scores or config.nlp_scores) _rouge_type = rouge_type or config.rouge_type - _register_score_udfs() + _register_score_udfs(schema_name) diff --git a/langkit/regexes.py b/langkit/regexes.py index a666a3c4..b0dc03eb 100644 --- a/langkit/regexes.py +++ b/langkit/regexes.py @@ -1,3 +1,4 @@ +from collections import defaultdict from copy import deepcopy from logging import getLogger @@ -6,7 +7,7 @@ from langkit import LangKitConfig, lang_config, prompt_column, response_column from whylogs.core.metrics.metrics import FrequentItemsMetric from whylogs.core.resolvers import MetricSpec -from typing import Optional, Set +from typing import Dict, Optional, Set from langkit.whylogs.unreg import unregister_udfs diagnostic_logger = getLogger(__name__) @@ -35,14 +36,17 @@ def wrappee(text): return wrappee -_registered: Set[str] = set() +_registered: Dict[str, Set[str]] = defaultdict( + set +) # _registered[schema_name] -> set of registered UDF names -def _register_udfs(config: Optional[LangKitConfig] = None): +def _register_udfs(config: Optional[LangKitConfig] = None, schema_name: str = ""): global _initialized _initialized = True global _registered - unregister_udfs(_registered) + unregister_udfs(_registered[schema_name], schema_name) + _registered[schema_name] = set() if config is None: config = lang_config default_metric_name = "has_patterns" @@ -56,8 +60,9 @@ def _register_udfs(config: Optional[LangKitConfig] = None): [column], udf_name=udf_name, metrics=[MetricSpec(FrequentItemsMetric)], + schema_name=schema_name, )(_wrapper(column, pattern_loader.get_regex_groups())) - _registered.add(udf_name) + _registered[schema_name].add(udf_name) if response_pattern_loader.get_regex_groups() is not None: column = response_column @@ -66,8 +71,9 @@ def _register_udfs(config: Optional[LangKitConfig] = None): [column], udf_name=udf_name, metrics=[MetricSpec(FrequentItemsMetric)], + schema_name=schema_name, )(_wrapper(column, response_pattern_loader.get_regex_groups())) - _registered.add(udf_name) + _registered[schema_name].add(udf_name) def init( @@ -75,6 +81,7 @@ def init( pattern_file_path: Optional[str] = None, config: Optional[LangKitConfig] = None, response_pattern_file_path: Optional[str] = None, + schema_name: str = "", ): global _initialized _initialized = True @@ -88,4 +95,4 @@ def init( pattern_loader = PatternLoader(config.pattern_file_path) response_pattern_loader = PatternLoader(config.response_pattern_file_path) - _register_udfs(config) + _register_udfs(config, schema_name) diff --git a/langkit/response_hallucination.py b/langkit/response_hallucination.py index 39a68700..1504fffc 100644 --- a/langkit/response_hallucination.py +++ b/langkit/response_hallucination.py @@ -1,12 +1,14 @@ +from collections import defaultdict from dataclasses import dataclass from logging import getLogger -from typing import List, Optional +from typing import List, Dict, Optional, Set from whylogs.experimental.core.udf_schema import register_dataset_udf from langkit import LangKitConfig, lang_config, prompt_column, response_column from nltk.tokenize import sent_tokenize from langkit.openai.openai import LLMInvocationParams, Conversation, ChatLog from langkit.transformer import Encoder from sentence_transformers import util +from langkit.whylogs.unreg import unregister_udfs diagnostic_logger = getLogger(__name__) @@ -258,14 +260,20 @@ def consistency_check(prompt: str, response: Optional[str] = None): checker: Optional[ConsistencyChecker] = None +_registered: Dict[str, Set[str]] = defaultdict( + set +) # _registered[schema_name] -> set of registered UDF names + + def init( language: Optional[str] = None, config: Optional[LangKitConfig] = None, llm: LLMInvocationParams = LLMInvocationParams(), num_samples=1, + schema_name: str = "", ): config = config or lang_config - global checker, embeddings_encoder + global checker, embeddings_encoder, _registered import nltk nltk.download("punkt") @@ -274,6 +282,10 @@ def init( ) embeddings_encoder = Encoder(config.response_transformer_name, custom_encoder=None) checker = ConsistencyChecker(llm, num_samples, embeddings_encoder) + unregister_udfs(_registered[schema_name], schema_name=schema_name) + _registered[schema_name] = set() + udf_name = f"{response_column}.hallucination" register_dataset_udf( - [prompt_column, response_column], f"{response_column}.hallucination" + [prompt_column, response_column], udf_name, schema_name=schema_name )(response_hallucination) + _registered[schema_name].add(udf_name) diff --git a/langkit/sentiment.py b/langkit/sentiment.py index d2a65542..037366ec 100644 --- a/langkit/sentiment.py +++ b/langkit/sentiment.py @@ -1,12 +1,15 @@ +from collections import defaultdict from copy import deepcopy -from typing import Optional, Set +from typing import Dict, Optional, Set from whylogs.experimental.core.udf_schema import register_dataset_udf from langkit import LangKitConfig, lang_config, prompt_column, response_column from langkit.whylogs.unreg import unregister_udfs -_registered: Set[str] = set() +_registered: Dict[str, Set[str]] = defaultdict( + set +) # _registered[schema_name] -> set of registered UDF names _nltk_downloaded = None @@ -118,12 +121,14 @@ def init( response_lexicon: Optional[str] = None, sentiment_model_path: Optional[str] = None, response_sentiment_model_path: Optional[str] = None, + schema_name: str = "", ): global _initialized _initialized = True global _registered - unregister_udfs(_registered) + unregister_udfs(_registered[schema_name], schema_name=schema_name) + _registered[schema_name] = set() config = config or deepcopy(lang_config) prompt_languages = {language} if language is not None else config.prompt_languages @@ -137,31 +142,39 @@ def init( if prompt_languages is not None and len(prompt_languages) > 0: if prompt_languages.issubset({"", "en"}) and _sentiment_analyzer: register_dataset_udf( - [prompt_column], udf_name=f"{prompt_column}.sentiment_nltk" + [prompt_column], + udf_name=f"{prompt_column}.sentiment_nltk", + schema_name=schema_name, )(_sentiment_wrapper(sentiment_nltk, _sentiment_analyzer, prompt_column)) - _registered.add(f"{prompt_column}.sentiment_nltk") + _registered[schema_name].add(f"{prompt_column}.sentiment_nltk") elif prompt_languages.issubset(_supported_languages) and _pipeline: register_dataset_udf( - [prompt_column], udf_name=f"{prompt_column}.sentiment_multi" + [prompt_column], + udf_name=f"{prompt_column}.sentiment_multi", + schema_name=schema_name, )(_sentiment_wrapper(sentiment_multilingual, _pipeline, prompt_column)) - _registered.add(f"{prompt_column}.sentiment_multi") + _registered[schema_name].add(f"{prompt_column}.sentiment_multi") if response_languages is not None and len(response_languages) > 0: if response_languages.issubset({"", "en"}) and _response_sentiment_analyzer: register_dataset_udf( - [response_column], udf_name=f"{response_column}.sentiment_nltk" + [response_column], + udf_name=f"{response_column}.sentiment_nltk", + schema_name=schema_name, )( _sentiment_wrapper( sentiment_nltk, _response_sentiment_analyzer, response_column ) ) - _registered.add(f"{response_column}.sentiment_nltk") + _registered[schema_name].add(f"{response_column}.sentiment_nltk") elif response_languages.issubset(_supported_languages) and _response_pipeline: register_dataset_udf( - [prompt_column], udf_name=f"{response_column}.sentiment_multi" + [prompt_column], + udf_name=f"{response_column}.sentiment_multi", + schema_name=schema_name, )( _sentiment_wrapper( sentiment_multilingual, _response_pipeline, prompt_column ) ) - _registered.add(f"{response_column}.sentiment_multi") + _registered[schema_name].add(f"{response_column}.sentiment_multi") diff --git a/langkit/textstat.py b/langkit/textstat.py index 977e8400..d35f91bc 100644 --- a/langkit/textstat.py +++ b/langkit/textstat.py @@ -1,3 +1,4 @@ +from collections import defaultdict from logging import getLogger from typing import Callable, Dict, List, Optional, Set, Tuple, Union from whylogs.core.stubs import pd @@ -72,10 +73,16 @@ def _unpack(t: Union[Tuple[str, str], Tuple[str, str, str]]) -> Tuple[str, str, return t if len(t) == 3 else (t[0], t[1], t[0]) # type: ignore -_registered: Set[str] = set() +_registered: Dict[str, Set[str]] = defaultdict( + set +) # _registered[schema_name] -> set of registered UDF names -def init(language: Optional[str] = None, config: Optional[LangKitConfig] = None): +def init( + language: Optional[str] = None, + config: Optional[LangKitConfig] = None, + schema_name: str = "", +): config = config or lang_config prompt_languages = ( {language} if language is not None else config.prompt_languages @@ -84,27 +91,31 @@ def init(language: Optional[str] = None, config: Optional[LangKitConfig] = None) {language} if language is not None else config.response_languages ) or set() global _registered - unregister_udfs(_registered) + unregister_udfs(_registered[schema_name], schema_name=schema_name) + _registered[schema_name] = set() for t in _udfs_to_register: - stat_name, schema_name, udf = _unpack(t) + stat_name, udf_lang, udf = _unpack(t) for column in [prompt_column, response_column]: - if schema_name in ( + if udf_lang in ( prompt_languages if column == prompt_column else response_languages ): udf_name = f"{column}.{udf}" register_dataset_udf( [column], udf_name=udf_name, - # schema_name=schema_name, # TODO: probably should be default schema + schema_name=schema_name, )(wrapper(stat_name, column)) - _registered.add(udf_name) + _registered[schema_name].add(udf_name) + for column in [prompt_column, response_column]: if "en" in ( prompt_languages if column == prompt_column else response_languages ): udf_name = f"{column}.aggregate_reading_level" - register_dataset_udf([column], udf_name=udf_name)(aggregate_wrapper(column)) - _registered.add(udf_name) + register_dataset_udf([column], udf_name=udf_name, schema_name=schema_name)( + aggregate_wrapper(column) + ) + _registered[schema_name].add(udf_name) diagnostic_logger.info("Initialized textstat metrics.") diff --git a/langkit/themes.py b/langkit/themes.py index 80ea7e22..7d9ef8fb 100644 --- a/langkit/themes.py +++ b/langkit/themes.py @@ -1,4 +1,5 @@ import json +from collections import defaultdict from copy import deepcopy from logging import getLogger from typing import Callable, Optional, Dict, List, Set @@ -59,12 +60,15 @@ def _map_embeddings(embeddings_map, theme_groups, transformer_model): ] -_registered: Set[str] = set() +_registered: Dict[str, Set[str]] = defaultdict( + set +) # _registered[schema_name] -> set of registered UDF names -def _register_theme_udfs(): +def _register_theme_udfs(schema_name: str): global _registered - unregister_udfs(_registered) + unregister_udfs(_registered[schema_name], schema_name=schema_name) + _registered[schema_name] = set() if _transformer_model is not None: _map_embeddings(_embeddings_map, _theme_groups, _transformer_model) for group in _theme_groups: @@ -72,8 +76,8 @@ def _register_theme_udfs(): if group == "refusal": continue udf_name = f"{column}.{group}_similarity" - _registered.add(udf_name) - register_dataset_udf([column], udf_name=udf_name)( + _registered[schema_name].add(udf_name) + register_dataset_udf([column], udf_name=udf_name, schema_name=schema_name)( create_similarity_function( group, column, _transformer_model, _embeddings_map ) @@ -90,8 +94,8 @@ def _register_theme_udfs(): if group == "jailbreak": continue udf_name = f"{column}.{group}_similarity" - _registered.add(udf_name) - register_dataset_udf([column], udf_name=udf_name)( + _registered[schema_name].add(udf_name) + register_dataset_udf([column], udf_name=udf_name, schema_name=schema_name)( create_similarity_function( group, column, @@ -128,6 +132,7 @@ def init( response_custom_encoder: Optional[Callable] = None, response_theme_file_path: Optional[str] = None, response_theme_json: Optional[str] = None, + schema_name: str = "", ): config = config or deepcopy(lang_config) global _transformer_model @@ -174,7 +179,7 @@ def init( else: _response_transformer_model = None - _register_theme_udfs() + _register_theme_udfs(schema_name) def get_subject_similarity( diff --git a/langkit/topics.py b/langkit/topics.py index fa549ec1..49717963 100644 --- a/langkit/topics.py +++ b/langkit/topics.py @@ -1,6 +1,7 @@ +from collections import defaultdict from copy import deepcopy from whylogs.experimental.core.udf_schema import register_dataset_udf -from typing import Callable, List, Optional, Set +from typing import Callable, Dict, List, Optional, Set from transformers import ( pipeline, ) @@ -41,7 +42,9 @@ def _wrapper(column: str, classifier, topics) -> Callable: return lambda text: [closest_topic(t, classifier, topics) for t in text[column]] -_registered: Set[str] = set() +_registered: Dict[str, Set[str]] = defaultdict( + set +) # _registered[schema_name] -> set of registered UDF names def init( @@ -53,11 +56,13 @@ def init( response_topics: Optional[List[str]] = None, response_model_path: Optional[str] = None, response_topic_classifier: Optional[str] = None, + schema_name: str = "", ): global _initialized _initialized = True global _registered - unregister_udfs(_registered) + unregister_udfs(_registered[schema_name], schema_name=schema_name) + _registered[schema_name] = set() config = config or deepcopy(lang_config) global _topics, _classifier _topics = topics or config.topics @@ -83,12 +88,16 @@ def init( if _classifier is not None: register_dataset_udf( - [prompt_column], udf_name=f"{prompt_column}.closest_topic" + [prompt_column], + udf_name=f"{prompt_column}.closest_topic", + schema_name=schema_name, )(_wrapper(prompt_column, _classifier, _topics)) - _registered.add(f"{prompt_column}.closest_topic") + _registered[schema_name].add(f"{prompt_column}.closest_topic") if _response_classifier is not None: register_dataset_udf( - [response_column], udf_name=f"{response_column}.closest_topic" + [response_column], + udf_name=f"{response_column}.closest_topic", + schema_name=schema_name, )(_wrapper(response_column, _response_classifier, _response_topics)) - _registered.add(f"{response_column}.closest_topic") + _registered[schema_name].add(f"{response_column}.closest_topic") diff --git a/langkit/toxicity.py b/langkit/toxicity.py index 22034ef6..0ef92a17 100644 --- a/langkit/toxicity.py +++ b/langkit/toxicity.py @@ -1,5 +1,6 @@ +from collections import defaultdict from copy import deepcopy -from typing import Optional, Set +from typing import Dict, Optional, Set from whylogs.experimental.core.udf_schema import register_dataset_udf from langkit import LangKitConfig, lang_config, prompt_column, response_column @@ -45,7 +46,9 @@ def _toxicity_wrapper(column, pipeline, tokenizer): return lambda text: [toxicity(t, pipeline, tokenizer) for t in text[column]] -_registered: Set[str] = set() +_registered: Dict[str, Set[str]] = defaultdict( + set +) # _registered[schema_name] -> set of registered UDF names def init( @@ -53,11 +56,13 @@ def init( model_path: Optional[str] = None, config: Optional[LangKitConfig] = None, response_model_path: Optional[str] = None, + schema_name: str = "", ): global _initialized _initialized = True global _registered - unregister_udfs(_registered) + unregister_udfs(_registered[schema_name], schema_name=schema_name) + _registered[schema_name] = set() from transformers import ( AutoModelForSequenceClassification, AutoTokenizer, @@ -79,14 +84,16 @@ def init( _toxicity_pipeline = TextClassificationPipeline( model=model, tokenizer=_toxicity_tokenizer, device=_device ) - register_dataset_udf([prompt_column], f"{prompt_column}.toxicity")( + register_dataset_udf( + [prompt_column], f"{prompt_column}.toxicity", schema_name=schema_name + )( translated_udf(translators)( _toxicity_wrapper( prompt_column, _toxicity_pipeline, _toxicity_tokenizer ) ) ) - _registered.add(f"{prompt_column}.toxicity") + _registered[schema_name].add(f"{prompt_column}.toxicity") model_path = response_model_path or config.response_toxicity_model_path global _response_toxicity_tokenizer, _response_toxicity_pipeline @@ -98,7 +105,9 @@ def init( _response_toxicity_pipeline = TextClassificationPipeline( model=model, tokenizer=_response_toxicity_tokenizer, device=_device ) - register_dataset_udf([response_column], f"{response_column}.toxicity")( + register_dataset_udf( + [response_column], f"{response_column}.toxicity", schema_name=schema_name + )( translated_udf(translators)( _toxicity_wrapper( response_column, @@ -107,4 +116,4 @@ def init( ) ) ) - _registered.add(f"{response_column}.toxicity") + _registered[schema_name].add(f"{response_column}.toxicity") diff --git a/langkit/whylogs/unreg.py b/langkit/whylogs/unreg.py index aa203347..2e229be2 100644 --- a/langkit/whylogs/unreg.py +++ b/langkit/whylogs/unreg.py @@ -20,8 +20,13 @@ def unregister_udf( if name in spec.udfs: found = True del spec.udfs[name] + if not found: logger.warn(f"UDF {name} could not be found for unregistering") + + us._multicolumn_udfs[schema_name] = [ + spec for spec in us._multicolumn_udfs[schema_name] if len(spec.udfs) > 0 + ] us._resolver_specs[schema_name] = list( filter(lambda x: x.column_name != name, us._resolver_specs[schema_name]) )