Upload 2 files

finetune.py

@@ -1,4 +1,5 @@
 import os
+import json
 
 CONTEXT_WINDOW = 1024 #has to fit in 4090
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -10,6 +11,7 @@ from transformers import (
 import torch
 from datasets import load_dataset
 from huggingface_hub import login
+from optimize_lr import best_lr
 
 # setup tokenizer
 tokenizer = AutoTokenizer.from_pretrained("Zyphra/Zamba2-1.2B-instruct", token=HF_TOKEN)
@@ -73,7 +75,7 @@ training_args = TrainingArguments(
     save_steps=500,
     save_total_limit=2,
     logging_steps=100,
-    learning_rate=2e-5,
+    learning_rate=best_lr,
     weight_decay=0.01,
     fp16=False,
     bf16=True,

optimize_lr.py

@@ -398,4 +398,7 @@ plot_gpr_results(study, final_optimization)
 
 # Save all results
 with open("lr_optimization_results.json", "w") as f:
-    json.dump(results, f, indent=4)
+    json.dump(results, f, indent=4)
+
+# Store best learning rate as a variable for finetune.py to use
+best_lr = study.best_params["learning_rate"]