Merge pull request #14 from MinishLab/add_push

add push_to_hub functionality

Author: stephantul

model2vec/model.py

@@ -2,6 +2,7 @@
 
 from logging import getLogger
 from pathlib import Path
+from tempfile import TemporaryDirectory
 from typing import Any, Iterator
 
 import numpy as np
@@ -10,7 +11,7 @@
 from torch.nn import EmbeddingBag
 from tqdm import tqdm
 
-from model2vec.utils import load_pretrained, save_pretrained
+from model2vec.utils import load_pretrained, push_folder_to_hub, save_pretrained
 
 PathLike = Path | str
 
@@ -154,3 +155,14 @@ def _encode_batch(self, sentences: list[str], max_length: int | None) -> np.ndar
     def _batch(sentences: list[str], batch_size: int) -> Iterator[list[str]]:
         """Batch the sentences into equal-sized."""
         return (sentences[i : i + batch_size] for i in range(0, len(sentences), batch_size))
+
+    def push_to_hub(self, repo_id: str, token: str | None) -> None:
+        """
+        Push the model to the huggingface hub.
+
+        :param repo_id: The repo id to push to.
+        :param token: The huggingface token to use.
+        """
+        with TemporaryDirectory() as temp_dir:
+            self.save_pretrained(temp_dir)
+            push_folder_to_hub(Path(temp_dir), repo_id, token)

model2vec/utils.py

@@ -100,3 +100,17 @@ def load_pretrained(
         )
 
     return embeddings, tokenizer, config
+
+
+def push_folder_to_hub(folder_path: Path, repo_id: str, huggingface_token: str | None) -> None:
+    """
+    Push a model folder to the huggingface hub.
+
+    :param folder_path: The path to the folder.
+    :param repo_id: The repo name.
+    :param huggingface_token: The huggingface token.
+    """
+    if not huggingface_hub.repo_exists(repo_id=repo_id, token=huggingface_token):
+        huggingface_hub.create_repo(repo_id, token=huggingface_token)
+    huggingface_hub.upload_folder(repo_id=repo_id, folder_path=folder_path, token=huggingface_token)
+    logger.info(f"Pushed model to {repo_id}")