Make token optional and private an argument, add template (#39)

* Fix modelcard comments

* Make token optional, add private as an optional argument which is set to False by default

* Update cite

* Resolved comment

---------

Co-authored-by: Pringled <thomas123@live.nl>

Author: stephantul

model2vec/model.py

@@ -209,15 +209,17 @@ 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:
+    def push_to_hub(self, repo_id: str, private: bool = False, token: str | None = None) -> None:
         """
         Push the model to the huggingface hub.
 
         NOTE: you need to pass a token if you are pushing a private model.
 
         :param repo_id: The repo id to push to.
+        :param private: Whether the repo, if created is set to private.
+            If the repo already exists, this doesn't change the visibility.
         :param token: The huggingface token to use.
         """
         with TemporaryDirectory() as temp_dir:
             self.save_pretrained(temp_dir, model_name=repo_id)
-            push_folder_to_hub(Path(temp_dir), repo_id, token)
+            push_folder_to_hub(Path(temp_dir), repo_id, private, token)

model2vec/model_card_template.md

@@ -4,9 +4,7 @@
 
 # {{ model_name }} Model Card
 
-Model2Vec distills a Sentence Transformer into a small, static model.
-This model is ideal for applications requiring fast, lightweight embeddings.
-
+This [Model2Vec](https://github.com/MinishLab/model2vec) model is a distilled version of {% if base_model %}the [{{ base_model }}](https://huggingface.co/{{ base_model }}){% else %}a{% endif %} Sentence Transformer. It uses static embeddings, allowing text embeddings to be computed orders of magnitude faster on both GPU and CPU. It is designed for applications where computational resources are limited or where real-time performance is critical.
 
 
 ## Installation
@@ -17,7 +15,7 @@ pip install model2vec
 ```
 
 ## Usage
-A StaticModel can be loaded using the `from_pretrained` method:
+Load this model using the `from_pretrained` method:
 ```python
 from model2vec import StaticModel
 
@@ -50,19 +48,20 @@ It works by passing a vocabulary through a sentence transformer model, then redu
 
 ## Additional Resources
 
+- [All Model2Vec models on the hub](https://huggingface.co/models?library=model2vec)
 - [Model2Vec Repo](https://github.com/MinishLab/model2vec)
 - [Model2Vec Results](https://github.com/MinishLab/model2vec?tab=readme-ov-file#results)
 - [Model2Vec Tutorials](https://github.com/MinishLab/model2vec/tree/main/tutorials)
 
-## Model Authors
+## Library Authors
 
 Model2Vec was developed by the [Minish Lab](https://github.com/MinishLab) team consisting of [Stephan Tulkens](https://github.com/stephantul) and [Thomas van Dongen](https://github.com/Pringled).
 
 ## Citation
 
 Please cite the [Model2Vec repository](https://github.com/MinishLab/model2vec) if you use this model in your work.
 ```
-@software{minishlab2024word2vec,
+@software{minishlab2024model2vec,
   authors = {Stephan Tulkens, Thomas van Dongen},
   title = {Model2Vec: Turn any Sentence Transformer into a Small Fast Model},
   year = {2024},

model2vec/utils.py

@@ -165,16 +165,17 @@ def load_pretrained(
     return embeddings, tokenizer, config
 
 
-def push_folder_to_hub(folder_path: Path, repo_id: str, token: str | None) -> None:
+def push_folder_to_hub(folder_path: Path, repo_id: str, private: bool, token: str | None) -> None:
     """
     Push a model folder to the huggingface hub, including model card.
 
     :param folder_path: The path to the folder.
     :param repo_id: The repo name.
+    :param private: Whether the repo is private.
     :param token: The huggingface token.
     """
     if not huggingface_hub.repo_exists(repo_id=repo_id, token=token):
-        huggingface_hub.create_repo(repo_id, token=token)
+        huggingface_hub.create_repo(repo_id, token=token, private=private)
 
     # Push model card and all model files to the Hugging Face hub
     huggingface_hub.upload_folder(repo_id=repo_id, folder_path=folder_path, token=token)