Upload folder using huggingface_hub

train/train_grpo.py

@@ -405,7 +405,18 @@ def main():
         processing_class=tokenizer,
     )
     trainer.train()
+    # CRITICAL: save_state() writes trainer_state.json to output_dir.
+    # save_model() alone only saves the adapter weights, NOT the per-step log.
+    # In Run #2, we missed save_state() and lost the reward history that drives the plot.
+    trainer.save_state()
     trainer.save_model(args.output_dir)
+    # Also explicitly write the log_history to a JSON we know we can find.
+    try:
+        log_path = Path(args.output_dir) / "log_history.json"
+        log_path.write_text(json.dumps(trainer.state.log_history, indent=2))
+        print(f"[done] log_history saved to {log_path} ({len(trainer.state.log_history)} entries)")
+    except Exception as e:
+        print(f"[warn] couldn't write log_history.json: {e}")
     print(f"[done] checkpoint saved to {args.output_dir}")
 
     # ----- HELD-OUT GENERALIZATION EVAL -----
@@ -415,6 +426,7 @@ def main():
     print(f"\n[held-out-eval] running trained model on {min(args.n_eval_clips, len(eval_ids))} held-out clips")
     eval_clip_ids = sorted(eval_ids)[: args.n_eval_clips]
     held_out_results = []
+    eval_failed = False
     model.eval()
     if hasattr(model, "gradient_checkpointing_disable"):
         try: model.gradient_checkpointing_disable()
@@ -422,46 +434,50 @@ def main():
     n_eval_correct = 0
     n_eval_well_formed = 0
     eval_rewards = []
-    for i, cid in enumerate(eval_clip_ids):
-        sc = scenarios[cid]
-        gold = "sarcastic" if sc["sarcasm"] else "sincere"
-        messages = [
-            {"role": "system", "content": SYSTEM_PROMPT},
-            {"role": "user", "content": build_full_observation(cid, scenarios)},
-        ]
-        encoded = tokenizer.apply_chat_template(
-            messages, return_tensors="pt", add_generation_prompt=True,
-        )
-        input_ids = encoded.input_ids if hasattr(encoded, "input_ids") else encoded
-        input_ids = input_ids.to(model.device)
-        prompt_len = input_ids.shape[1]
-        with _t.no_grad():
-            out = model.generate(
-                input_ids=input_ids,
-                max_new_tokens=args.max_completion_length,
-                do_sample=False,
-                pad_token_id=tokenizer.eos_token_id,
-                use_cache=True,
+    try:
+        for i, cid in enumerate(eval_clip_ids):
+            sc = scenarios[cid]
+            gold = "sarcastic" if sc["sarcasm"] else "sincere"
+            messages = [
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": build_full_observation(cid, scenarios)},
+            ]
+            encoded = tokenizer.apply_chat_template(
+                messages, return_tensors="pt", add_generation_prompt=True,
             )
-        text = tokenizer.decode(out[0][prompt_len:], skip_special_tokens=True)
-        decomp = reward_decomposition(text, gold)
-        held_out_results.append({
-            "clip_id": cid,
-            "gold": gold,
-            "is_pivot": bool(sc.get("is_pivot")),
-            "predicted": decomp["_predicted"],
-            "confidence": decomp["_confidence"],
-            "correct": decomp["_correct"],
-            "well_formed": decomp["_well_formed"],
-            "reward_total": decomp["_total"],
-            "completion_text": text[:1500],
-        })
-        eval_rewards.append(decomp["_total"])
-        if decomp["_correct"]: n_eval_correct += 1
-        if decomp["_well_formed"]: n_eval_well_formed += 1
-        if (i + 1) % 20 == 0:
-            print(f"   [{i+1}/{len(eval_clip_ids)}] running mean reward = {sum(eval_rewards)/len(eval_rewards):.3f}, "
-                  f"correct so far = {n_eval_correct}/{i+1}", flush=True)
+            input_ids = encoded.input_ids if hasattr(encoded, "input_ids") else encoded
+            input_ids = input_ids.to(model.device)
+            prompt_len = input_ids.shape[1]
+            with _t.no_grad():
+                out = model.generate(
+                    input_ids=input_ids,
+                    max_new_tokens=args.max_completion_length,
+                    do_sample=False,
+                    pad_token_id=tokenizer.eos_token_id,
+                    use_cache=True,
+                )
+            text = tokenizer.decode(out[0][prompt_len:], skip_special_tokens=True)
+            decomp = reward_decomposition(text, gold)
+            held_out_results.append({
+                "clip_id": cid,
+                "gold": gold,
+                "is_pivot": bool(sc.get("is_pivot")),
+                "predicted": decomp["_predicted"],
+                "confidence": decomp["_confidence"],
+                "correct": decomp["_correct"],
+                "well_formed": decomp["_well_formed"],
+                "reward_total": decomp["_total"],
+                "completion_text": text[:1500],
+            })
+            eval_rewards.append(decomp["_total"])
+            if decomp["_correct"]: n_eval_correct += 1
+            if decomp["_well_formed"]: n_eval_well_formed += 1
+            if (i + 1) % 20 == 0:
+                print(f"   [{i+1}/{len(eval_clip_ids)}] running mean reward = {sum(eval_rewards)/len(eval_rewards):.3f}, "
+                      f"correct so far = {n_eval_correct}/{i+1}", flush=True)
+    except Exception as e:
+        print(f"[error] held-out eval crashed at clip {i}: {e}")
+        eval_failed = True
     eval_summary = {
         "n_eval_clips": len(eval_clip_ids),
         "mean_reward": sum(eval_rewards) / max(1, len(eval_rewards)),
@@ -498,36 +514,38 @@ def main():
         except Exception as e:
             print(f"[error] push_to_hub failed: {e}")
 
-    # Push trainer_state.json (with per-step rewards/loss) to HF Space for plotting
+    # Push trainer_state.json + log_history.json + held_out_eval.json to HF Space.
+    # Each upload is wrapped individually so a partial network failure doesn't
+    # kill the whole script. We need at least the held_out_eval JSON to land.
     if args.save_trainer_state_to_hub_space:
-        try:
-            from huggingface_hub import HfApi
-            from pathlib import Path as _P
-            state_path = _P(args.output_dir) / "trainer_state.json"
-            if state_path.exists():
-                HfApi().upload_file(
-                    path_or_fileobj=str(state_path),
-                    path_in_repo="data/trainer_state_run1.json",
-                    repo_id=args.save_trainer_state_to_hub_space,
+        from huggingface_hub import HfApi
+        from pathlib import Path as _P
+        repo_id = args.save_trainer_state_to_hub_space
+        run_tag = _P(args.output_dir).name  # e.g. "run3"
+        api = HfApi()
+        for local_name, hub_name, label in [
+            ("trainer_state.json",   f"data/trainer_state_{run_tag}.json",   "trainer_state"),
+            ("log_history.json",     f"data/log_history_{run_tag}.json",     "log_history"),
+            ("held_out_eval.json",   f"data/held_out_eval_{run_tag}.json",   "held_out_eval"),
+        ]:
+            path = _P(args.output_dir) / local_name
+            if not path.exists():
+                print(f"[warn] {local_name} not found at {path}, skipping upload")
+                continue
+            try:
+                api.upload_file(
+                    path_or_fileobj=str(path),
+                    path_in_repo=hub_name,
+                    repo_id=repo_id,
                     repo_type="space",
-                    commit_message=f"GRPO Run #1 trainer_state ({args.max_steps} steps)",
+                    commit_message=f"GRPO {run_tag} {label} ({args.max_steps} steps)",
                 )
-                print(f"[done] trainer_state.json pushed to {args.save_trainer_state_to_hub_space}/data/trainer_state_run1.json")
-            else:
-                print(f"[warn] trainer_state.json not found at {state_path}")
-            # Also push held_out_eval.json (the proof of generalization)
-            held_out_path = _P(args.output_dir) / "held_out_eval.json"
-            if held_out_path.exists():
-                HfApi().upload_file(
-                    path_or_fileobj=str(held_out_path),
-                    path_in_repo="data/held_out_eval_run1.json",
-                    repo_id=args.save_trainer_state_to_hub_space,
-                    repo_type="space",
-                    commit_message=f"GRPO Run #1 held-out eval ({args.n_eval_clips} clips)",
-                )
-                print(f"[done] held_out_eval.json pushed")
-        except Exception as e:
-            print(f"[error] save_trainer_state_to_hub_space failed: {e}")
+                print(f"[done] {label} pushed to {repo_id}/{hub_name}")
+            except Exception as e:
+                print(f"[error] upload {label} failed: {e}")
+
+    print(f"\n[main] subtext-arena GRPO run finished cleanly.")
+    sys.exit(0)
 
 
 if __name__ == "__main__":