feat: add pushing the model to hf

jina-ai · Aug 25, 2023 · 09d2811 · 09d2811
1 parent 693deb2
commit 09d2811
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/textbook/train.py b/textbook/train.py
@@ -49,6 +49,7 @@ def train(
     wandb_log_model: Optional[
         bool
     ] = None,  # will be true by default if use_wandb is true
+    push_model_to_hf: bool = False, # if set, will push the model to hf
     local_rank: Annotated[int, typer.Option("--local_rank")] = 0,
     deepspeed: Optional[str] = None,
     debug: bool = False,
@@ -127,12 +128,13 @@ def train(
 
     trainer.train()
 
-    model.save_pretrained("jinaai/starcoder-1b-textbook")
-    tokenizer.save_pretrained("jinaai/starcoder-1b-textbook")
+    if push_model_to_hf:
+        model.save_pretrained(output_dir)
+        tokenizer.save_pretrained(output_dir)
 
-    # Push to the hub
-    model.push_to_hub("jinaai/starcoder-1b-textbook")
-    tokenizer.push_to_hub("jinaai/starcoder-1b-textbook")
+        # Push to the hub
+        model.push_to_hub('jinaai/starcoder-1b-textbook')
+        tokenizer.push_to_hub('jinaai/starcoder-1b-textbook')
 
     accuracy_results, sample_results = evaluate(
         model, tokenizer, eval_size=eval_size, max_new_tokens=eval_max_new_tokens