Upload train_grpo.py with huggingface_hub

train_grpo.py

@@ -5,6 +5,7 @@
 #     "peft>=0.7.0",
 #     "transformers>=4.45.0",
 #     "accelerate>=0.24.0",
+#     "huggingface_hub>=0.20.0",
 #     "trackio",
 #     "datasets",
 #     "bitsandbytes",
@@ -23,11 +24,13 @@ Usage:
     uv run train_grpo.py --sft-model tobil/qmd-query-expansion-0.6B
 """
 
+import os
 import re
 import torch
 import trackio
 from collections import Counter
 from datasets import load_dataset
+from huggingface_hub import login
 from peft import LoraConfig, PeftModel, get_peft_model
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from trl import GRPOTrainer, GRPOConfig
@@ -256,6 +259,14 @@ def main():
         print(f"  LR: {args.lr}")
         return
 
+    # Login to HuggingFace Hub
+    hf_token = os.environ.get("HF_TOKEN")
+    if hf_token:
+        print("Logging in to HuggingFace Hub...")
+        login(token=hf_token)
+    else:
+        print("Warning: HF_TOKEN not set, will try cached login")
+
     # Load dataset (just prompts needed for GRPO)
     print("Loading dataset...")
     dataset = load_dataset("tobil/qmd-query-expansion-train", split="train")