Feat/llm description scorer (#236)

* decompose `DescriptionScorer` into two scorers

* first version of llm description scorer

* add concurrency

* tweak prompt a little bit

* fix imports

* fix typing

* upd docstring to llm description scorer

* refactor description scorer tests

* fix stupid bugs

* upd optimization test config

* implement generator dumper

* fix type errors

* skip LLMScorer test if openai is not available

* fix presets test

* fix config test

* upd inference test

Author: voorhs

autointent/_dump_tools/__init__.py

@@ -1,4 +1,3 @@
 from .main import Dumper
-from .unit_dumpers import PydanticModelDumper
 
-__all__ = ["Dumper", "PydanticModelDumper"]
+__all__ = ["Dumper"]

autointent/_dump_tools/generator_dumper.py

@@ -0,0 +1,24 @@
+"""Separate file to fix circular import error."""
+
+from pathlib import Path
+from typing import Any
+
+from autointent.generation import Generator
+
+from .base import BaseObjectDumper
+
+
+class GeneratorDumper(BaseObjectDumper[Generator]):
+    dir_or_file_name = "generators"
+
+    @staticmethod
+    def dump(obj: Generator, path: Path, exists_ok: bool) -> None:
+        obj.dump(path, exist_ok=exists_ok)
+
+    @staticmethod
+    def load(path: Path, **kwargs: Any) -> Generator:  # noqa: ANN401, ARG004
+        return Generator.load(path)
+
+    @classmethod
+    def check_isinstance(cls, obj: Any) -> bool:  # noqa: ANN401
+        return isinstance(obj, Generator)

autointent/_dump_tools/main.py

@@ -10,6 +10,7 @@
 from autointent.schemas import TagsList
 
 from .base import BaseObjectDumper, ModuleAttributes, ModuleSimpleAttributes
+from .generator_dumper import GeneratorDumper
 from .unit_dumpers import (
     ArraysDumper,
     CatBoostDumper,
@@ -46,6 +47,7 @@ class Dumper:
         HFTokenizerDumper,
         TorchModelDumper,
         CatBoostDumper,
+        GeneratorDumper,
     ]
 
     @staticmethod

autointent/_presets/heavy.yaml

@@ -13,7 +13,12 @@ search_space:
         k:
           low: 1
           high: 20
-      - module_name: description
+      - module_name: description_bi
+        temperature:
+          low: 0.01
+          high: 10
+          log: true
+      - module_name: description_cross
         temperature:
           low: 0.01
           high: 10

autointent/generation/__init__.py

@@ -3,6 +3,7 @@
 See :ref:`data-aug-tuts`.
 """
 
-from ._generator import Generator
+from ._cache import StructuredOutputCache
+from ._generator import Generator, RetriesExceededError
 
-__all__ = ["Generator"]
+__all__ = ["Generator", "RetriesExceededError"]

autointent/generation/_cache.py

@@ -9,7 +9,7 @@
 from dotenv import load_dotenv
 from pydantic import BaseModel, ValidationError
 
-from autointent._dump_tools import PydanticModelDumper
+from autointent._dump_tools.unit_dumpers import PydanticModelDumper
 from autointent._hash import Hasher
 from autointent.generation.chat_templates import Message
 

autointent/generation/_generator.py

@@ -1,9 +1,11 @@
 """Wrapper class for accessing OpenAI API."""
 
+import json
 import logging
 import os
+from pathlib import Path
 from textwrap import dedent
-from typing import Any, ClassVar, Literal, TypeVar
+from typing import Any, ClassVar, Literal, TypedDict, TypeVar
 
 import openai
 from dotenv import load_dotenv
@@ -22,6 +24,27 @@
 """Type variable for Pydantic models used in structured output generation."""
 
 
+class GeneratorDumpData(TypedDict):
+    use_cache: bool
+    model_name: str
+    base_url: str | None
+    generation_params: dict[str, Any]
+
+
+class RetriesExceededError(RuntimeError):
+    """Exception raised when LLM call fails after all retry attempts."""
+
+    def __init__(self, max_retries: int, messages: list[Message]) -> None:
+        """Initialize the error with retry count and messages.
+
+        Args:
+            max_retries: Maximum number of retry attempts that were made
+            messages: Messages that were sent to the LLM
+        """
+        msg = f"LLM call failed after {max_retries + 1} attempts. Messages: {messages}"
+        super().__init__(msg)
+
+
 class Generator:
     """Wrapper class for accessing OpenAI API.
 
@@ -32,6 +55,8 @@ class Generator:
         **generation_params: kwargs that will be sent with a request to the endpoint.
     """
 
+    _dump_data_filename = "init_params.json"
+
     _default_generation_params: ClassVar[dict[str, Any]] = {
         "max_tokens": 150,
         "n": 1,
@@ -57,17 +82,23 @@ def __init__(
         """
         base_url = base_url or os.getenv("OPENAI_BASE_URL")
         model_name = model_name or os.getenv("OPENAI_MODEL_NAME")
+
         if model_name is None:
             msg = "Specify model_name arg or OPENAI_MODEL_NAME environment variable"
             raise ValueError(msg)
+
         self.model_name = model_name
+        self.base_url = base_url
+        self.use_cache = use_cache
+
         self.client = openai.OpenAI(base_url=base_url)
         self.async_client = openai.AsyncOpenAI(base_url=base_url)
+        self.cache = StructuredOutputCache(use_cache=use_cache)
+
         self.generation_params = {
             **self._default_generation_params,
             **generation_params,
         }  #  https://stackoverflow.com/a/65539348
-        self.cache = StructuredOutputCache(use_cache=use_cache)
 
     def get_chat_completion(self, messages: list[Message]) -> str:
         """Prompt LLM and return its answer.
@@ -221,12 +252,8 @@ async def get_structured_output_async(
             current_messages.extend(self._create_retry_messages(error, raw))
 
         if res is None:
-            msg = (
-                f"Failed to generate valid structured output after {max_retries + 1} attempts.\n"
-                f"Messages: {current_messages}"
-            )
             logger.exception(msg)
-            raise RuntimeError(msg)
+            raise RetriesExceededError(max_retries=max_retries, messages=current_messages)
 
         # Cache the successful result
         self.cache.set(messages, output_model, backend, self.generation_params, res)
@@ -338,14 +365,32 @@ def get_structured_output_sync(
             current_messages.extend(self._create_retry_messages(error, raw))
 
         if res is None:
-            msg = (
-                f"Failed to generate valid structured output after {max_retries + 1} attempts.\n"
-                f"Messages: {current_messages}"
-            )
             logger.exception(msg)
-            raise RuntimeError(msg)
+            raise RetriesExceededError(max_retries=max_retries, messages=current_messages)
 
         # Cache the successful result
         self.cache.set(messages, output_model, backend, self.generation_params, res)
 
         return res
+
+    def dump(self, path: Path, exist_ok: bool = True) -> None:
+        data: GeneratorDumpData = {
+            "base_url": self.base_url,
+            "generation_params": self.generation_params,
+            "model_name": self.model_name,
+            "use_cache": self.use_cache,
+        }
+
+        path.mkdir(exist_ok=exist_ok, parents=True)
+
+        with (path / self._dump_data_filename).open("w", encoding="utf-8") as file:
+            json.dump(data, file, indent=4, ensure_ascii=False)
+
+    @classmethod
+    def load(cls, path: Path) -> "Generator":
+        with (path / cls._dump_data_filename).open(encoding="utf-8") as file:
+            data: GeneratorDumpData = json.load(file)
+
+        generation_params = data.pop("generation_params")  # type: ignore[misc]
+
+        return cls(**data, **generation_params)

autointent/modules/__init__.py

@@ -15,12 +15,14 @@
 from .scoring import (
     BERTLoRAScorer,
     BertScorer,
+    BiEncoderDescriptionScorer,
     CatBoostScorer,
     CNNScorer,
-    DescriptionScorer,
+    CrossEncoderDescriptionScorer,
     DNNCScorer,
     KNNScorer,
     LinearScorer,
+    LLMDescriptionScorer,
     MLKnnScorer,
     PTuningScorer,
     RerankScorer,
@@ -47,7 +49,9 @@ def _create_modules_dict(modules: list[type[T]]) -> dict[str, type[T]]:
         DNNCScorer,
         KNNScorer,
         LinearScorer,
-        DescriptionScorer,
+        BiEncoderDescriptionScorer,
+        CrossEncoderDescriptionScorer,
+        LLMDescriptionScorer,
         RerankScorer,
         SklearnScorer,
         MLKnnScorer,
@@ -62,28 +66,3 @@ def _create_modules_dict(modules: list[type[T]]) -> dict[str, type[T]]:
 DECISION_MODULES: dict[str, type[BaseDecision]] = _create_modules_dict(
     [ArgmaxDecision, JinoosDecision, ThresholdDecision, TunableDecision, AdaptiveDecision],
 )
-
-
-__all__ = [
-    "AdaptiveDecision",
-    "ArgmaxDecision",
-    "BaseDecision",
-    "BaseEmbedding",
-    "BaseModule",
-    "BaseRegex",
-    "BaseScorer",
-    "CatBoostScorer",
-    "DNNCScorer",
-    "DescriptionScorer",
-    "JinoosDecision",
-    "KNNScorer",
-    "LinearScorer",
-    "LogregAimedEmbedding",
-    "MLKnnScorer",
-    "RerankScorer",
-    "RetrievalAimedEmbedding",
-    "SimpleRegex",
-    "SklearnScorer",
-    "ThresholdDecision",
-    "TunableDecision",
-]

autointent/modules/scoring/__init__.py

@@ -1,6 +1,6 @@
 from ._bert import BertScorer
 from ._catboost import CatBoostScorer
-from ._description import DescriptionScorer
+from ._description import BiEncoderDescriptionScorer, CrossEncoderDescriptionScorer, LLMDescriptionScorer
 from ._dnnc import DNNCScorer
 from ._knn import KNNScorer, RerankScorer
 from ._linear import LinearScorer
@@ -13,11 +13,13 @@
 __all__ = [
     "BERTLoRAScorer",
     "BertScorer",
+    "BiEncoderDescriptionScorer",
     "CNNScorer",
     "CatBoostScorer",
+    "CrossEncoderDescriptionScorer",
     "DNNCScorer",
-    "DescriptionScorer",
     "KNNScorer",
+    "LLMDescriptionScorer",
     "LinearScorer",
     "MLKnnScorer",
     "PTuningScorer",

autointent/modules/scoring/_description/__init__.py

@@ -1,3 +1,5 @@
-from .description import DescriptionScorer
+from .bi_encoder import BiEncoderDescriptionScorer
+from .cross_encoder import CrossEncoderDescriptionScorer
+from .llm_encoder import LLMDescriptionScorer
 
-__all__ = ["DescriptionScorer"]
+__all__ = ["BiEncoderDescriptionScorer", "CrossEncoderDescriptionScorer", "LLMDescriptionScorer"]