diff --git a/.bumpversion.cfg b/.bumpversion.cfg
index 06c3aab..0311aba 100644
--- a/.bumpversion.cfg
+++ b/.bumpversion.cfg
@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.0.28.dev1
+current_version = 0.0.28.dev2
 tag = False
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(\.(?P<release>[a-z]+)(?P<build>\d+))?
 serialize = 
diff --git a/langkit/metrics/library.py b/langkit/metrics/library.py
index 14b55cc..0f5bef6 100644
--- a/langkit/metrics/library.py
+++ b/langkit/metrics/library.py
@@ -113,14 +113,19 @@ def __call__(self) -> MetricCreator:
                 return self.toxicity_score()
 
             @staticmethod
-            def toxicity_score() -> MetricCreator:
+            def toxicity_score(onnx: bool = True) -> MetricCreator:
                 """
                 Analyze the input for toxicity. The output of this metric ranges from 0 to 1, where 0 indicates
                 non-toxic and 1 indicates toxic.
                 """
-                from langkit.metrics.toxicity_onnx import prompt_toxicity_metric
+                if onnx:
+                    from langkit.metrics.toxicity_onnx import prompt_toxicity_metric
+
+                    return prompt_toxicity_metric
+                else:
+                    from langkit.metrics.toxicity import prompt_toxicity_metric
 
-                return prompt_toxicity_metric
+                    return prompt_toxicity_metric
 
         class stats:
             def __call__(self) -> MetricCreator:
@@ -331,14 +336,19 @@ def __call__(self) -> MetricCreator:
                 return self.toxicity_score()
 
             @staticmethod
-            def toxicity_score() -> MetricCreator:
+            def toxicity_score(onnx: bool = True) -> MetricCreator:
                 """
                 Analyze the toxicity of the response. The output of this metric ranges from 0 to 1, where 0
                 indicates a non-toxic response and 1 indicates a toxic response.
                 """
-                from langkit.metrics.toxicity_onnx import response_toxicity_metric
+                if onnx:
+                    from langkit.metrics.toxicity_onnx import response_toxicity_metric
+
+                    return response_toxicity_metric
+                else:
+                    from langkit.metrics.toxicity import response_toxicity_metric
 
-                return response_toxicity_metric
+                    return response_toxicity_metric
 
         class stats:
             def __call__(self) -> MetricCreator:
diff --git a/pyproject.toml b/pyproject.toml
index a403a37..c0c08f3 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langkit"
-version = "0.0.28.dev1"
+version = "0.0.28.dev2"
 description = "A language toolkit for monitoring LLM interactions"
 authors = ["WhyLabs.ai <langkit@whylabs.ai>"]
 homepage = "https://docs.whylabs.ai/docs/large-language-model-monitoring"