add get_default_hfmodel_config

Author: voorhs

src/autointent/_pipeline/_pipeline.py

@@ -22,6 +22,7 @@
     LoggingConfig,
     VectorIndexConfig,
     get_default_embedder_config,
+    get_default_hfmodel_config,
     get_default_vector_index_config,
 )
 from autointent.custom_types import NodeType
@@ -64,7 +65,7 @@ def __init__(
             self.embedder_config = get_default_embedder_config()
             self.cross_encoder_config = CrossEncoderConfig()
             self.data_config = DataConfig()
-            self.transformer_config = HFModelConfig()
+            self.transformer_config = get_default_hfmodel_config()
             self.hpo_config = HPOConfig()
             self.vector_index_config = get_default_vector_index_config()
         elif not isinstance(nodes[0], InferenceNode):

src/autointent/configs/__init__.py

@@ -17,6 +17,7 @@
     EmbedderFineTuningConfig,
     HFModelConfig,
     TokenizerConfig,
+    get_default_hfmodel_config,
 )
 from ._vector_index import FaissConfig, OpenSearchConfig, VectorIndexConfig, get_default_vector_index_config
 
@@ -40,6 +41,7 @@
     "VectorIndexConfig",
     "VocabConfig",
     "get_default_embedder_config",
+    "get_default_hfmodel_config",
     "get_default_vector_index_config",
     "initialize_embedder_config",
 ]

src/autointent/configs/_transformers.py

@@ -53,6 +53,7 @@ class HFModelConfig(BaseModel):
     fp16: bool = Field(False, description="Whether to use mixed precision training (not all devices support this).")
     tokenizer_config: TokenizerConfig = Field(default_factory=TokenizerConfig)
     trust_remote_code: bool = Field(False, description="Whether to trust the remote code when loading the model.")
+    revision: str | None = Field(None, description="Revision from HF repo")
 
     @classmethod
     def from_search_config(cls, values: dict[str, Any] | str | BaseModel | None) -> Self:
@@ -75,6 +76,10 @@ def from_search_config(cls, values: dict[str, Any] | str | BaseModel | None) ->
         return cls(**values)
 
 
+def get_default_hfmodel_config() -> HFModelConfig:
+    return HFModelConfig(model_name="prajjwal1/bert-tiny", revision="refs/pr/16")
+
+
 class CrossEncoderConfig(HFModelConfig):
     model_name: str = Field("cross-encoder/ms-marco-MiniLM-L6-v2", description="Name of the hugging face model.")
     train_head: bool = Field(

src/autointent/context/_context.py

@@ -16,6 +16,7 @@
     HPOConfig,
     LoggingConfig,
     VectorIndexConfig,
+    get_default_hfmodel_config,
 )
 
 from .data_handler import DataHandler
@@ -25,9 +26,7 @@
     from pathlib import Path
 
     from autointent import Dataset
-    from autointent.configs import (
-        DataConfig,
-    )
+    from autointent.configs import DataConfig
 
 
 class Context:
@@ -202,4 +201,4 @@ def resolve_transformer(self) -> HFModelConfig:
         """
         if hasattr(self, "transformer_config"):
             return self.transformer_config
-        return HFModelConfig()
+        return get_default_hfmodel_config()