Add checkpoint download

README.md

@@ -123,6 +123,8 @@ Use the code below to get started with the model.
 from utils.regression_head import RegressionHead
 from transformers import AutoModel, AutoTokenizer
 from torch import bfloat16, no_grad, nn
+from huggingface_hub import hf_hub_download
+
 
 embedding_model = AutoModel.from_pretrained(
     'Snowflake/snowflake-arctic-embed-m-v2.0', 
@@ -132,8 +134,14 @@ embedding_model = AutoModel.from_pretrained(
     add_pooling_layer=False,
 ).cuda()
 
+file_path = hf_hub_download(
+    repo_id="Jackal-AI/JQL-Edu-Heads",
+    filename="checkpoints/edu-mistral-snowflake-balanced.ckpt",
+    repo_type="model"
+)
+
 regression_head = RegressionHead.load_from_checkpoint(
-    '../models/gemma-snowflake-balanced_wandb_run_id=uqw4xkyx.ckpt', 
+    file_path, 
     map_location='cuda'
 ).to(bfloat16)