Update src/litbench_pipeline.py

src/litbench_pipeline.py

@@ -32,8 +32,13 @@ import time
 import torch
 from peft.peft_model import PeftModel
 from datasets import load_dataset
+from huggingface_hub import login
 
-
+# Load token from Hugging Face Secrets
+HF_TOKEN = os.getenv("HF_TOKEN")
+# Login using the token (only needed if loading private models)
+if HF_TOKEN:
+    login(HF_TOKEN)
 
 # Function to determine the chatbot's first message based on user choices
 def setup(download_option, train_option):