feat: replace 8m by 32m for training

MinishLab · Feb 11, 2025 · acd24d7 · acd24d7
1 parent 370f0ab
commit acd24d7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/model2vec/train/base.py b/model2vec/train/base.py
@@ -45,7 +45,7 @@ def construct_head(self) -> nn.Sequential:
 
     @classmethod
     def from_pretrained(
-        cls: type[ModelType], out_dim: int = 2, model_name: str = "minishlab/potion-base-8m", **kwargs: Any
+        cls: type[ModelType], out_dim: int = 2, model_name: str = "minishlab/potion-base-32m", **kwargs: Any
     ) -> ModelType:
         """Load the model from a pretrained model2vec model."""
         model = StaticModel.from_pretrained(model_name)