implement new hashing strategy

Author: voorhs

autointent/_embedder.py

@@ -10,6 +10,7 @@
 from pathlib import Path
 from typing import TypedDict
 
+import huggingface_hub
 import numpy as np
 import numpy.typing as npt
 import torch
@@ -20,7 +21,7 @@
 from .configs import EmbedderConfig, TaskTypeEnum
 
 
-def get_embeddings_path(filename: str) -> Path:
+def _get_embeddings_path(filename: str) -> Path:
     """Get the path to the embeddings file.
 
     This function constructs the full path to an embeddings file stored
@@ -37,6 +38,18 @@ def get_embeddings_path(filename: str) -> Path:
     return Path(user_cache_dir("autointent")) / "embeddings" / f"{filename}.npy"
 
 
+def _get_latest_commit_hash(model_name: str) -> str:
+    """Get the latest commit hash for a given Hugging Face model.
+
+    Args:
+        model_name: The name of the model to get the latest commit hash for.
+
+    Returns:
+        The latest commit hash for the given model name.
+    """
+    return huggingface_hub.model_info(model_name, revision="main").sha
+
+
 class EmbedderDumpMetadata(TypedDict):
     """Metadata for saving and loading an Embedder instance."""
 
@@ -63,7 +76,6 @@ class Embedder:
 
     _metadata_dict_name: str = "metadata.json"
     _dump_dir: Path | None = None
-    config: EmbedderConfig
     embedding_model: SentenceTransformer
 
     def __init__(self, embedder_config: EmbedderConfig) -> None:
@@ -74,14 +86,6 @@ def __init__(self, embedder_config: EmbedderConfig) -> None:
         """
         self.config = embedder_config
 
-        self.embedding_model = SentenceTransformer(
-            self.config.model_name,
-            device=self.config.device,
-            prompts=embedder_config.get_prompt_config(),
-            similarity_fn_name=self.config.similarity_fn_name,
-            trust_remote_code=self.config.trust_remote_code,
-        )
-
         self._logger = logging.getLogger(__name__)
 
     def __hash__(self) -> int:
@@ -91,11 +95,27 @@ def __hash__(self) -> int:
             The hash value of the Embedder.
         """
         hasher = Hasher()
-        for parameter in self.embedding_model.parameters():
-            hasher.update(parameter.detach().cpu().numpy())
+        if self.config.freeze:
+            commit_hash = _get_latest_commit_hash(self.config.model_name)
+            hasher.update(commit_hash)
+        else:
+            self._load_model()
+            for parameter in self.embedding_model.parameters():
+                hasher.update(parameter.detach().cpu().numpy())
         hasher.update(self.config.tokenizer_config.max_length)
         return hasher.intdigest()
 
+    def _load_model(self) -> None:
+        """Load sentence transformers model to device."""
+        if not hasattr(self, "embedding_model"):
+            self.embedding_model = SentenceTransformer(
+                self.config.model_name,
+                device=self.config.device,
+                prompts=self.config.get_prompt_config(),
+                similarity_fn_name=self.config.similarity_fn_name,
+                trust_remote_code=self.config.trust_remote_code,
+            )
+
     def clear_ram(self) -> None:
         """Move the embedding model to CPU and delete it from memory."""
         self._logger.debug("Clearing embedder %s from memory", self.config.model_name)
@@ -165,10 +185,12 @@ def embed(self, utterances: list[str], task_type: TaskTypeEnum | None = None) ->
             hasher.update(self)
             hasher.update(utterances)
 
-            embeddings_path = get_embeddings_path(hasher.hexdigest())
+            embeddings_path = _get_embeddings_path(hasher.hexdigest())
             if embeddings_path.exists():
                 return np.load(embeddings_path)  # type: ignore[no-any-return]
 
+        self._load_model()
+
         self._logger.debug(
             "Calculating embeddings with model %s, batch_size=%d, max_seq_length=%s, embedder_device=%s",
             self.config.model_name,

autointent/configs/_transformers.py

@@ -64,6 +64,8 @@ class EmbedderConfig(HFModelConfig):
     similarity_fn_name: str | None = Field(
         "cosine", description="Name of the similarity function to use (cosine, dot, euclidean, manhattan)."
     )
+    use_cache: bool = Field(True, description="Whether to use embeddings caching.")
+    freeze: bool = Field(True, description="Whether to freeze the model parameters.")
 
     def get_prompt_config(self) -> dict[str, str] | None:
         """Get the prompt config for the given prompt type.
@@ -111,7 +113,6 @@ def get_prompt_type(self, prompt_type: TaskTypeEnum | None) -> str | None:  # no
             return self.default_prompt
         assert_never(prompt_type)
 
-    use_cache: bool = Field(False, description="Whether to use embeddings caching.")
 
 
 class CrossEncoderConfig(HFModelConfig):