feat(scoring): TypeSafe (jev) zero-shot description scorer

.github/workflows/ci.yaml

@@ -138,7 +138,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        dependency-group: [ "base", "transformers", "peft", "catboost" ]
+        dependency-group: [ "base", "transformers", "peft", "catboost", "typesafe" ]
     uses: ./.github/workflows/reusable-test.yaml
     secrets: inherit
     with:

docs/source/installation.rst

@@ -45,6 +45,9 @@ The table below lists each published extra, what it is typically used for in Aut
    * - ``openai``
      - OpenAI API clients for OpenAI-backed embedders and the ``zero-shot-llm`` style workflows.
      - ``pip install "autointent[openai]"``
+   * - ``typesafe``
+     - TypeSafe API client for the ``description_typesafe`` scorer and the ``zero-shot-typesafe`` preset.
+     - ``pip install "autointent[typesafe]"``
    * - ``dspy``
      - DSPy-based augmentation and related generation utilities.
      - ``pip install "autointent[dspy]"``

docs/source/quickstart.rst

@@ -190,6 +190,7 @@ Available Presets
 - ``transformers-light``: Transformer models with limited search
 - ``transformers-no-hpo``: Transformer models without hyperparameter optimization
 - ``zero-shot-llm``: Zero-shot classification using OpenAI models
+- ``zero-shot-typesafe``: Zero-shot classification using TypeSafe's ``jev`` model (requires ``TYPESAFE_API_KEY``)
 - ``zero-shot-encoders``: Zero-shot classification using transformer models
 
 Making Predictions

pyproject.toml

@@ -80,6 +80,9 @@ openai = [
     "openai (>=2,<3)",
     "tiktoken (>=0.7,<1)",
 ]
+typesafe = [
+    "typesafe-sdk (>=0.7,<1)",
+]
 vllm = [
     "vllm>=0.20.0",
 ]
@@ -316,6 +319,7 @@ module = [
     "dspy.evaluate.auto_evaluation",
     "codecarbon",
     "catboost",
+    "typesafe_sdk",
     "jupytext",
     "docs_utils.*",
 ]

src/autointent/_deps.py

@@ -37,6 +37,7 @@
     "peft",
     "sentence-transformers",
     "transformers",
+    "typesafe",
     "vllm",
     "wandb",
 ]

src/autointent/_dump_tools/unit_dumpers.py

@@ -3,7 +3,7 @@
 import importlib
 import json
 import logging
-from typing import TYPE_CHECKING, Any, TypeVar
+from typing import TYPE_CHECKING, Any, TypeVar, cast
 
 import aiofiles
 import joblib
@@ -215,7 +215,8 @@ def dump(obj: PeftModel, path: Path, exists_ok: bool) -> None:
             ptuning_path = path / "ptuning"
             ptuning_path.mkdir(parents=True, exist_ok=exists_ok)
             obj.save_pretrained(str(ptuning_path / "peft"))
-            obj.base_model.save_pretrained(ptuning_path / "base_model")
+            base_model = cast("PreTrainedModel", obj.base_model)
+            base_model.save_pretrained(ptuning_path / "base_model")
         else:
             # strategy to save lora models: merge adapters and save as usual hugging face model
             lora_path = path / "lora"

src/autointent/_presets/zero-shot-typesafe.yaml

@@ -0,0 +1,24 @@
+search_space:
+  - node_type: scoring
+    target_metric: scoring_f1
+    search_space:
+      - module_name: description_typesafe
+        question_type: [choice, noul]
+        temperature:
+          low: 0.3
+          high: 3.0
+  - node_type: decision
+    target_metric: decision_accuracy
+    search_space:
+      - module_name: threshold
+        thresh:
+          low: 0.1
+          high: 0.9
+      - module_name: argmax
+      - module_name: jinoos
+      - module_name: tunable
+      - module_name: adaptive
+hpo_config:
+  sampler: tpe
+  n_trials: 30
+  n_startup_trials: 10

src/autointent/advisor/_estimates/_resource.py

@@ -106,13 +106,14 @@ def _embedder_model_name(embedder: EmbedderConfig) -> str | None:
         "description_bi",
         "description_cross",
         "description_llm",
+        "description_typesafe",
     },
 )
 
 # Cache-honoring modules whose per-entry estimate already bundles the embedder
 # forward into `time_hours` (vs. classic linear/catboost which don't).
 _EMBEDDER_FORWARD_TRANSFORMER_MODULES = frozenset(
-    {"knn", "mlknn", "retrieval", "description_bi", "description_cross", "description_llm"},
+    {"knn", "mlknn", "retrieval", "description_bi", "description_cross", "description_llm", "description_typesafe"},
 )
 
 
@@ -608,6 +609,7 @@ def _emit_resource_findings(
         "description_bi",
         "description_cross",
         "description_llm",
+        "description_typesafe",
     },
 )
 

src/autointent/advisor/_runner.py

@@ -274,7 +274,7 @@ def _data_phase(
                 )
 
     # partial descriptions x description scorer
-    description_modules = {"description_bi", "description_cross", "description_llm"}
+    description_modules = {"description_bi", "description_cross", "description_llm", "description_typesafe"}
     has_description = any(e.get("module_name") in description_modules for _, e in _walk_modules(search_space))
     if has_description and stats.has_descriptions is False:
         report.add(

src/autointent/advisor/_workflows.py

@@ -45,6 +45,7 @@
     "transformers-light",
     "nn-heavy",
     "zero-shot-llm",
+    "zero-shot-typesafe",
     "nn-medium",
     "classic-heavy",
     "transformers-no-hpo",

src/autointent/custom_types/_types.py

@@ -126,6 +126,7 @@ class Split:
     "transformers-light",
     "transformers-no-hpo",
     "zero-shot-llm",
+    "zero-shot-typesafe",
     "zero-shot-encoders",
 ]
 """Bundled search-space presets that our library supports.