Update finetune.py

finetune.py

@@ -17,13 +17,9 @@ from peft import (
 
 HF_TOKEN = os.environ.get("TRL_TOKEN", None)
 if HF_TOKEN:
-    try:
-        shutil.rmtree("./data/")
-    except:
-        pass
 
     repo = Repository(
-        local_dir="./data/", clone_from="gustavoaq/llama_ft", use_auth_token=HF_TOKEN, repo_type="models"
+        local_dir="./checkpoints/", clone_from="gustavoaq/llama_ft", use_auth_token=HF_TOKEN, repo_type="models"
     )
     repo.git_pull()
 # Parameters
@@ -163,6 +159,8 @@ trainer = transformers.Trainer(
         save_total_limit=100,
         load_best_model_at_end=True if VAL_SET_SIZE > 0 else False,
         ddp_find_unused_parameters=False if ddp else None,
+        push_to_hub=True,
+        push_to_hub_model_id='llama_ft'
         
     ),
     data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),
@@ -175,6 +173,8 @@ model.state_dict = (
 ).__get__(model, type(model))
 
 trainer.train()
-trainer.push_to_hub()
 
 model.save_pretrained(OUTPUT_DIR)
+trainer.push_to_hub()
+
+