Refactor/another revision (#73)

* change default `clear_ram` value

* fix some references

* convert search space guide to ipynb

* fix some references

* rename tutorials, fix refs, begin with python api tutorial

* stage progress on python api tutorial

* move device option to embedder config

* stage progress on tutorials

Author: voorhs

autointent/_embedder.py

@@ -158,7 +158,7 @@ def embed(self, utterances: list[str]) -> npt.NDArray[np.float32]:
                 return np.load(embeddings_path)  # type: ignore[no-any-return]
 
         self.logger.debug(
-            "Calculating embeddings with model %s, batch_size=%d, max_seq_length=%s, device=%s",
+            "Calculating embeddings with model %s, batch_size=%d, max_seq_length=%s, embedder_device=%s",
             self.model_name,
             self.batch_size,
             str(self.max_length),

autointent/configs/_optimization_cli.py

@@ -43,7 +43,7 @@ class LoggingConfig:
     """Path to the directory where the modules will be dumped. If None, the modules will not be dumped"""
     dump_modules: bool = False
     """Whether to dump the modules or not"""
-    clear_ram: bool = True
+    clear_ram: bool = False
     """Whether to clear the RAM after dumping the modules"""
 
     def __post_init__(self) -> None:
@@ -77,8 +77,6 @@ class VectorIndexConfig:
 
     db_dir: Path | None = None
     """Path to the directory where the vector index database will be saved. If None, the database will not be saved"""
-    device: str = "cpu"
-    """Device to use for the vector index. Can be 'cpu', 'cuda', 'cuda:0', 'mps', etc."""
     save_db: bool = False
     """Whether to save the vector index database or not"""
 
@@ -109,6 +107,8 @@ class EmbedderConfig:
     """Max length for the embedder. If None, the max length will be taken from model config"""
     use_cache: bool = False
     """Flag indicating whether to cache embeddings for reuse, improving performance in repeated operations."""
+    device: str = "cpu"
+    """Device to use for the vector index. Can be 'cpu', 'cuda', 'cuda:0', 'mps', etc."""
 
 
 @dataclass

autointent/context/_context.py

@@ -65,7 +65,7 @@ def configure_vector_index(self, config: VectorIndexConfig, embedder_config: Emb
         self.embedder_config = embedder_config
 
         self.vector_index_client = VectorIndexClient(
-            self.vector_index_config.device,
+            self.embedder_config.device,
             self.vector_index_config.db_dir,
             self.embedder_config.batch_size,
             self.embedder_config.max_length,
@@ -115,7 +115,7 @@ def get_inference_config(self) -> dict[str, Any]:
         nodes_configs = [asdict(cfg) for cfg in self.optimization_info.get_inference_nodes_config()]
         return {
             "metadata": {
-                "device": self.get_device(),
+                "embedder_device": self.get_device(),
                 "multilabel": self.is_multilabel(),
                 "n_classes": self.get_n_classes(),
                 "seed": self.seed,
@@ -168,11 +168,11 @@ def get_db_dir(self) -> Path:
 
     def get_device(self) -> str:
         """
-        Get the device used by the vector index client.
+        Get the embedder device used by the vector index client.
 
         :return: Device name.
         """
-        return self.vector_index_client.device
+        return self.vector_index_client.embedder_device
 
     def get_batch_size(self) -> int:
         """

autointent/context/vector_index_client/_vector_index.py

@@ -28,7 +28,7 @@ class VectorIndex:
     def __init__(
         self,
         model_name: str,
-        device: str,
+        embedder_device: str,
         embedder_batch_size: int = 32,
         embedder_max_length: int | None = None,
         embedder_use_cache: bool = False,
@@ -37,7 +37,7 @@ def __init__(
         Initialize the vector index.
 
         :param model_name: Name of the embedding model to use.
-        :param device: Device for running the embedding model (e.g., "cpu", "cuda").
+        :param embedder_device: Device for running the embedding model (e.g., "cpu", "cuda").
         :param embedder_batch_size: Batch size for the embedder.
         :param embedder_max_length: Maximum sequence length for the embedder.
         :param embedder_use_cache: Flag indicating whether to cache intermediate embeddings.
@@ -46,11 +46,11 @@ def __init__(
         self.embedder = Embedder(
             model_name=model_name,
             batch_size=embedder_batch_size,
-            device=device,
+            device=embedder_device,
             max_length=embedder_max_length,
             use_cache=embedder_use_cache,
         )
-        self.device = device
+        self.embedder_device = embedder_device
 
         self.labels: list[LabelType] = []  # (n_samples,) or (n_samples, n_classes)
         self.texts: list[str] = []
@@ -200,7 +200,7 @@ def load(self, dir_path: Path) -> None:
         """
         self.dump_dir = Path(dir_path)
         self.index = faiss.read_index(str(dir_path / "index.faiss"))
-        self.embedder = Embedder(model_name=dir_path / "embedding_model", device=self.device)
+        self.embedder = Embedder(model_name=dir_path / "embedding_model", device=self.embedder_device)
         with (dir_path / "texts.json").open() as file:
             self.texts = json.load(file)
         with (dir_path / "labels.json").open() as file:

autointent/context/vector_index_client/_vector_index_client.py

@@ -28,7 +28,7 @@ class VectorIndexClient:
 
     def __init__(
         self,
-        device: str,
+        embedder_device: str,
         db_dir: str | Path | None,
         embedder_batch_size: int = 32,
         embedder_max_length: int | None = None,
@@ -37,14 +37,14 @@ def __init__(
         """
         Initialize the VectorIndexClient.
 
-        :param device: Device to run the embedding model on.
+        :param embedder_device: Device to run the embedding model on.
         :param db_dir: Directory for storing vector indexes. Defaults to a cache directory.
         :param embedder_batch_size: Batch size for the embedding model.
         :param embedder_max_length: Maximum sequence length for the embedding model.
         :param embedder_use_cache: Flag indicating whether to cache intermediate embeddings.
         """
         self._logger = logging.getLogger(__name__)
-        self.device = device
+        self.embedder_device = embedder_device
         self.db_dir = get_db_dir(db_dir)
         self.embedder_batch_size = embedder_batch_size
         self.embedder_max_length = embedder_max_length
@@ -69,7 +69,7 @@ def create_index(
 
         index = VectorIndex(
             model_name,
-            self.device,
+            self.embedder_device,
             self.embedder_batch_size,
             self.embedder_max_length,
             self.embedder_use_cache,
@@ -176,7 +176,7 @@ def get_index(self, model_name: str) -> VectorIndex:
         if dirpath is not None:
             index = VectorIndex(
                 model_name,
-                self.device,
+                self.embedder_device,
                 self.embedder_batch_size,
                 self.embedder_max_length,
                 self.embedder_use_cache,

autointent/modules/retrieval/_vectordb.py

@@ -69,7 +69,7 @@ def __init__(
         k: int,
         embedder_name: str,
         db_dir: str | None = None,
-        device: str = "cpu",
+        embedder_device: str = "cpu",
         batch_size: int = 32,
         max_length: int | None = None,
         embedder_use_cache: bool = False,
@@ -80,13 +80,13 @@ def __init__(
         :param k: Number of nearest neighbors to retrieve.
         :param embedder_name: Name of the embedder used for creating embeddings.
         :param db_dir: Path to the database directory. If None, defaults will be used.
-        :param device: Device to run operations on, e.g., "cpu" or "cuda".
+        :param embedder_device: Device to run operations on, e.g., "cpu" or "cuda".
         :param batch_size: Batch size for embedding generation.
         :param max_length: Maximum sequence length for embeddings. None if not set.
         :param embedder_use_cache: Flag indicating whether to cache intermediate embeddings.
         """
         self.embedder_name = embedder_name
-        self.device = device
+        self.embedder_device = embedder_device
         self._db_dir = db_dir
         self.batch_size = batch_size
         self.max_length = max_length
@@ -113,7 +113,7 @@ def from_context(
             k=k,
             embedder_name=embedder_name,
             db_dir=str(context.get_db_dir()),
-            device=context.get_device(),
+            embedder_device=context.get_device(),
             batch_size=context.get_batch_size(),
             max_length=context.get_max_length(),
             embedder_use_cache=context.get_use_cache(),
@@ -138,7 +138,7 @@ def fit(self, utterances: list[str], labels: list[LabelType]) -> None:
         :param labels: List of corresponding labels for the utterances.
         """
         vector_index_client = VectorIndexClient(
-            self.device,
+            self.embedder_device,
             self.db_dir,
             embedder_batch_size=self.batch_size,
             embedder_max_length=self.max_length,
@@ -212,7 +212,7 @@ def load(self, path: str) -> None:
             self.metadata: VectorDBMetadata = json.load(file)
 
         vector_index_client = VectorIndexClient(
-            device=self.device,
+            embedder_device=self.embedder_device,
             db_dir=self.metadata["db_dir"],
             embedder_batch_size=self.metadata["batch_size"],
             embedder_max_length=self.metadata["max_length"],

autointent/modules/scoring/_description/description.py

@@ -54,7 +54,7 @@ def __init__(
         self,
         embedder_name: str,
         temperature: float = 1.0,
-        device: str = "cpu",
+        embedder_device: str = "cpu",
         batch_size: int = 32,
         max_length: int | None = None,
         embedder_use_cache: bool = False,
@@ -64,13 +64,13 @@ def __init__(
 
         :param embedder_name: Name of the embedder model.
         :param temperature: Temperature parameter for scaling logits, defaults to 1.0.
-        :param device: Device to run the embedder on, e.g., "cpu" or "cuda".
+        :param embedder_device: Device to run the embedder on, e.g., "cpu" or "cuda".
         :param batch_size: Batch size for embedding generation, defaults to 32.
         :param max_length: Maximum sequence length for embedding, defaults to None.
         :param embedder_use_cache: Flag indicating whether to cache intermediate embeddings.
         """
         self.temperature = temperature
-        self.device = device
+        self.embedder_device = embedder_device
         self.embedder_name = embedder_name
         self.batch_size = batch_size
         self.max_length = max_length
@@ -99,7 +99,7 @@ def from_context(
 
         instance = cls(
             temperature=temperature,
-            device=context.get_device(),
+            embedder_device=context.get_device(),
             embedder_name=embedder_name,
             embedder_use_cache=context.get_use_cache(),
         )
@@ -139,7 +139,7 @@ def fit(
         if self.precomputed_embeddings:
             # this happens only when LinearScorer is within Pipeline opimization after RetrievalNode optimization
             vector_index_client = VectorIndexClient(
-                self.device,
+                self.embedder_device,
                 self.db_dir,
                 self.batch_size,
                 self.max_length,
@@ -153,7 +153,7 @@ def fit(
             embedder = vector_index.embedder
         else:
             embedder = Embedder(
-                device=self.device,
+                device=self.embedder_device,
                 model_name=self.embedder_name,
                 batch_size=self.batch_size,
                 max_length=self.max_length,
@@ -230,7 +230,7 @@ def load(self, path: str) -> None:
 
         embedder_dir = dump_dir / self.embedding_model_subdir
         self.embedder = Embedder(
-            device=self.device,
+            device=self.embedder_device,
             model_name=embedder_dir,
             batch_size=self.metadata["batch_size"],
             max_length=self.metadata["max_length"],

autointent/modules/scoring/_dnnc/dnnc.py

@@ -313,7 +313,7 @@ def load(self, path: str) -> None:
         self.n_classes = self.metadata["n_classes"]
 
         vector_index_client = VectorIndexClient(
-            device=self.device,
+            embedder_device=self.device,
             db_dir=self.metadata["db_dir"],
             embedder_batch_size=self.metadata["batch_size"],
             embedder_max_length=self.metadata["max_length"],

autointent/modules/scoring/_knn/knn.py

@@ -90,7 +90,7 @@ def __init__(
         k: int,
         weights: WEIGHT_TYPES = "distance",
         db_dir: str | None = None,
-        device: str = "cpu",
+        embedder_device: str = "cpu",
         batch_size: int = 32,
         max_length: int | None = None,
         embedder_use_cache: bool = False,
@@ -105,7 +105,7 @@ def __init__(
             - "distance" (or True): Weight inversely proportional to distance.
             - "closest": Only the closest neighbor of each class is weighted.
         :param db_dir: Path to the database directory, or None to use default.
-        :param device: Device to run operations on, e.g., "cpu" or "cuda".
+        :param embedder_device: Device to run operations on, e.g., "cpu" or "cuda".
         :param batch_size: Batch size for embedding generation, defaults to 32.
         :param max_length: Maximum sequence length for embedding, or None for default.
         :param embedder_use_cache: Flag indicating whether to cache intermediate embeddings.
@@ -114,7 +114,7 @@ def __init__(
         self.k = k
         self.weights = weights
         self._db_dir = db_dir
-        self.device = device
+        self.embedder_device = embedder_device
         self.batch_size = batch_size
         self.max_length = max_length
         self.embedder_use_cache = embedder_use_cache
@@ -158,7 +158,7 @@ def from_context(
             k=k,
             weights=weights,
             db_dir=str(context.get_db_dir()),
-            device=context.get_device(),
+            embedder_device=context.get_device(),
             batch_size=context.get_batch_size(),
             max_length=context.get_max_length(),
             embedder_use_cache=context.get_use_cache(),
@@ -188,7 +188,9 @@ def fit(self, utterances: list[str], labels: list[LabelType]) -> None:
         else:
             self.n_classes = len(set(labels))
             self.multilabel = False
-        vector_index_client = VectorIndexClient(self.device, self.db_dir, embedder_use_cache=self.embedder_use_cache)
+        vector_index_client = VectorIndexClient(
+            self.embedder_device, self.db_dir, embedder_use_cache=self.embedder_use_cache
+        )
 
         if self.prebuilt_index:
             # this happens only after RetrievalNode optimization
@@ -265,7 +267,7 @@ def _restore_state_from_metadata(self, metadata: KNNScorerDumpMetadata) -> None:
         self.multilabel = metadata["multilabel"]
 
         vector_index_client = VectorIndexClient(
-            device=self.device,
+            embedder_device=self.embedder_device,
             db_dir=metadata["db_dir"],
             embedder_batch_size=metadata["batch_size"],
             embedder_max_length=metadata["max_length"],

autointent/modules/scoring/_knn/rerank_scorer.py

@@ -51,7 +51,7 @@ def __init__(
         m: int | None = None,
         rank_threshold_cutoff: int | None = None,
         db_dir: str | None = None,
-        device: str = "cpu",
+        embedder_device: str = "cpu",
         batch_size: int = 32,
         max_length: int | None = None,
     ) -> None:
@@ -68,7 +68,7 @@ def __init__(
         :param m: Number of top-ranked neighbors to consider, or None to use k.
         :param rank_threshold_cutoff: Rank threshold cutoff for re-ranking, or None.
         :param db_dir: Path to the database directory, or None to use default.
-        :param device: Device to run operations on, e.g., "cpu" or "cuda".
+        :param embedder_device: Device to run operations on, e.g., "cpu" or "cuda".
         :param batch_size: Batch size for embedding generation, defaults to 32.
         :param max_length: Maximum sequence length for embedding, or None for default.
         """
@@ -77,7 +77,7 @@ def __init__(
             k=k,
             weights=weights,
             db_dir=db_dir,
-            device=device,
+            embedder_device=embedder_device,
             batch_size=batch_size,
             max_length=max_length,
         )
@@ -123,7 +123,7 @@ def from_context(
             m=m,
             rank_threshold_cutoff=rank_threshold_cutoff,
             db_dir=str(context.get_db_dir()),
-            device=context.get_device(),
+            embedder_device=context.get_device(),
             batch_size=context.get_batch_size(),
             max_length=context.get_max_length(),
         )
@@ -138,7 +138,7 @@ def fit(self, utterances: list[str], labels: list[LabelType]) -> None:
         :param utterances: List of utterances to fit the scorer.
         :param labels: List of labels corresponding to the utterances.
         """
-        self._scorer = CrossEncoder(self.cross_encoder_name, device=self.device, max_length=self.max_length)  # type: ignore[arg-type]
+        self._scorer = CrossEncoder(self.cross_encoder_name, device=self.embedder_device, max_length=self.max_length)  # type: ignore[arg-type]
 
         super().fit(utterances, labels)
 
@@ -179,7 +179,7 @@ def _restore_state_from_metadata(self, metadata: RerankScorerDumpMetadata) -> No
         self.m = metadata["m"] if metadata["m"] else self.k
         self.cross_encoder_name = metadata["cross_encoder_name"]
         self.rank_threshold_cutoff = metadata["rank_threshold_cutoff"]
-        self._scorer = CrossEncoder(self.cross_encoder_name, device=self.device, max_length=self.max_length)  # type: ignore[arg-type]
+        self._scorer = CrossEncoder(self.cross_encoder_name, device=self.embedder_device, max_length=self.max_length)  # type: ignore[arg-type]
 
     def _predict(self, utterances: list[str]) -> tuple[npt.NDArray[Any], list[list[str]]]:
         """