Uploading models to hub

src/config.py

@@ -7,9 +7,11 @@ MAX_DOWNLOAD_TIME = 0.2
 
 IMAGE_DOWNLOAD_PATH = pathlib.Path("./data/images")
 WANDB_LOG_PATH = pathlib.Path("/tmp/wandb_logs")
+MODEL_PATH = pathlib.Path("/tmp/models")
 
 IMAGE_DOWNLOAD_PATH.mkdir(parents=True, exist_ok=True)
 WANDB_LOG_PATH.mkdir(parents=True, exist_ok=True)
+MODEL_PATH.mkdir(parents=True, exist_ok=True)
 
 MODEL_NAME = "tiny_clip"
 

src/trainer.py

@@ -1,3 +1,5 @@
+import os
+
 from src import config
 from src import data
 from src import loss
@@ -8,7 +10,28 @@ from src import utils
 from src.lightning_module import LightningModule
 
 
+def _upload_model_to_hub(
+    vision_encoder: models.TinyCLIPVisionEncoder, text_encoder: models.TinyCLIPTextEncoder
+):
+    vision_encoder.save_pretrained(
+        str(config.MODEL_PATH),
+        variant="vision_encoder",
+        safe_serialization=True,
+        push_to_hub=True,
+        repo_id="debug-clip-model",
+    )
+    text_encoder.save_pretrained(
+        str(config.MODEL_PATH),
+        variant="text_encoder",
+        safe_serialization=True,
+        push_to_hub=True,
+        repo_id="debug-clip-model",
+    )
+
+
 def train(trainer_config: config.TrainerConfig):
+    if "HF_TOKEN" not in os.environ:
+        raise ValueError("Please set the HF_TOKEN environment variable.")
     transform = vision_model.get_vision_transform(trainer_config._model_config.vision_config)
     tokenizer = tk.Tokenizer(trainer_config._model_config.text_config)
     train_dl, valid_dl = data.get_dataset(
@@ -28,6 +51,8 @@ def train(trainer_config: config.TrainerConfig):
     trainer = utils.get_trainer(trainer_config)
     trainer.fit(lightning_module, train_dl, valid_dl)
 
+    _upload_model_to_hub(vision_encoder, text_encoder)
+
 
 if __name__ == "__main__":
     trainer_config = config.TrainerConfig(debug=True)