Merge branch 'dev' into feat/augmentation

Author: Darinochka

.github/workflows/generate-schema.yaml

@@ -0,0 +1,45 @@
+name: Generate JSON Schema
+
+on:
+  push:
+    branches:
+      - main
+
+permissions:
+  contents: write
+
+jobs:
+  generate-schema:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v4
+
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.10'
+
+    - name: Install dependencies
+      run: |
+          pip install .
+
+    - name: Generate JSON Schema
+      run: python scripts/generate_json_schema_config.py
+
+    - name: Check for changes
+      id: check_changes
+      run: |
+        git diff --exit-code docs/optimizer_config.schema.json || echo "changed=true" >> $GITHUB_ENV
+
+    - name: Commit and push changes
+      if: env.changed == 'true'
+      env:
+        GITHUB_TOKEN: ${{ github.token }}
+      run: |
+        git config --global user.name "github-actions[bot]"
+        git config --global user.email "github-actions[bot]@users.noreply.github.com"
+        git add docs/optimizer_config.schema.json
+        git commit -m "Update optimizer_config.schema.json"
+        git push

.vscode/settings.json

@@ -2,5 +2,11 @@
     "ruff.configuration": "pyproject.toml",
     "python.analysis.extraPaths": [
         "./docs/source"
-    ]
+    ],
+    "yaml.schemas": {
+        "./docs/optimizer_config.schema.json": [
+            "*.yaml",
+            "!*/.github/*/*.yaml"
+        ]
+    }
 }

autointent/_datafiles/default-multiclass-config.yaml

@@ -1,27 +1,26 @@
 # TODO: make up a better and more versatile config
 - node_type: embedding
-  metric: retrieval_hit_rate
+  target_metric: retrieval_hit_rate
   search_space:
     - module_name: retrieval
       k: [10]
       embedder_name:
         - avsolatorio/GIST-small-Embedding-v0
-        - infgrad/stella-base-en-v2
+        - sergeyzh/rubert-tiny-turbo
 - node_type: scoring
-  metric: scoring_roc_auc
+  target_metric: scoring_roc_auc
   search_space:
     - module_name: knn
       k: [1, 3, 5, 10]
       weights: ["uniform", "distance", "closest"]
     - module_name: linear
     - module_name: dnnc
       cross_encoder_name:
-        - BAAI/bge-reranker-base
         - cross-encoder/ms-marco-MiniLM-L-6-v2
       k: [1, 3, 5, 10]
 - node_type: decision
-  metric: decision_accuracy
+  target_metric: decision_accuracy
   search_space:
     - module_name: threshold
       thresh: [0.5]
-    - module_name: argmax
+    - module_name: argmax

autointent/_datafiles/default-multilabel-config.yaml

@@ -1,21 +1,21 @@
 # TODO: make up a better and more versatile config
 - node_type: embedding
-  metric: retrieval_hit_rate_intersecting
+  target_metric: retrieval_hit_rate_intersecting
   search_space:
     - module_name: retrieval
       k: [10]
       embedder_name:
         - deepvk/USER-bge-m3
 - node_type: scoring
-  metric: scoring_roc_auc
+  target_metric: scoring_roc_auc
   search_space:
     - module_name: knn
       k: [3]
       weights: ["uniform", "distance", "closest"]
     - module_name: linear
 - node_type: decision
-  metric: decision_accuracy
+  target_metric: decision_accuracy
   search_space:
     - module_name: threshold
       thresh: [0.5]
-    - module_name: adaptive
+    - module_name: adaptive

autointent/_datafiles/inference-config-example.yaml

@@ -2,7 +2,7 @@
   module_name: retrieval
   module_config:
     k: 10
-    model_name: infgrad/stella-base-en-v2
+    model_name: sergeyzh/rubert-tiny-turbo
   load_path: .
 - node_type: scoring
   module_name: knn

autointent/_dataset/_dataset.py

@@ -100,13 +100,14 @@ def from_hub(cls, repo_id: str) -> "Dataset":
         :param repo_id: ID of the Hugging Face repository.
         :return: Initialized Dataset object.
         """
-        splits, intents = load_dataset(repo_id), []
+        from ._reader import DictReader
+
+        splits = load_dataset(repo_id)
+        mapping = dict(**splits)
         if Split.INTENTS in get_dataset_config_names(repo_id):
-            intents = load_dataset(repo_id, Split.INTENTS)[Split.INTENTS].to_list()
-        return cls(
-            splits.items(),
-            intents=[Intent.model_validate(intent) for intent in intents],
-        )
+            mapping["intents"] = load_dataset(repo_id, Split.INTENTS)[Split.INTENTS].to_list()
+
+        return DictReader().read(mapping)
 
     def to_multilabel(self) -> "Dataset":
         """

autointent/_pipeline/_pipeline.py

@@ -13,6 +13,7 @@
 from autointent.custom_types import ListOfGenericLabels, NodeType
 from autointent.metrics import PREDICTION_METRICS_MULTILABEL
 from autointent.nodes import InferenceNode, NodeOptimizer
+from autointent.nodes.schemes import OptimizationConfig
 from autointent.utils import load_default_search_space, load_search_space
 
 from ._schemas import InferencePipelineOutput, InferencePipelineUtteranceOutput
@@ -72,10 +73,12 @@ def from_search_space(cls, search_space: list[dict[str, Any]] | Path | str, seed
         Create pipeline optimizer from dictionary search space.
 
         :param search_space: Dictionary config
+        :param seed: random seed
         """
         if isinstance(search_space, Path | str):
             search_space = load_search_space(search_space)
-        nodes = [NodeOptimizer(**node) for node in search_space]
+        validated_search_space = OptimizationConfig(search_space).model_dump()  # type: ignore[arg-type]
+        nodes = [NodeOptimizer(**node) for node in validated_search_space]
         return cls(nodes=nodes, seed=seed)
 
     @classmethod
@@ -84,6 +87,9 @@ def default_optimizer(cls, multilabel: bool, seed: int = 42) -> "Pipeline":
         Create pipeline optimizer with default search space for given classification task.
 
         :param multilabel: Whether the task multi-label, or single-label.
+        :param seed: random seed
+
+        :return: Pipeline
         """
         return cls.from_search_space(search_space=load_default_search_space(multilabel), seed=seed)
 

autointent/custom_types.py

@@ -5,7 +5,7 @@
 """
 
 from enum import Enum
-from typing import Literal, TypeAlias, TypedDict
+from typing import Literal, TypeAlias
 
 
 class LogLevel(Enum):
@@ -46,10 +46,6 @@ class LogLevel(Enum):
 """
 
 
-class BaseMetadataDict(TypedDict):
-    """Base metadata dictionary for storing additional information."""
-
-
 class NodeType(str, Enum):
     """Enumeration of node types in the AutoIntent pipeline."""
 

autointent/generation/utterances/__init__.py

@@ -0,0 +1,14 @@
+from .basic import SynthesizerChatTemplate, UtteranceGenerator
+from .evolution import AbstractEvolution, ConcreteEvolution, EvolutionChatTemplate, ReasoningEvolution, UtteranceEvolver
+from .generator import Generator
+
+__all__ = [
+    "AbstractEvolution",
+    "ConcreteEvolution",
+    "EvolutionChatTemplate",
+    "Generator",
+    "ReasoningEvolution",
+    "SynthesizerChatTemplate",
+    "UtteranceEvolver",
+    "UtteranceGenerator",
+]

autointent/generation/utterances/basic/__init__.py

@@ -0,0 +1,4 @@
+from .chat_template import SynthesizerChatTemplate
+from .utterance_generator import UtteranceGenerator
+
+__all__ = ["SynthesizerChatTemplate", "UtteranceGenerator"]