From 24a533c985918cd04d1758dd19e22933f9b7bb25 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 5 Feb 2025 05:36:04 -0800
Subject: [PATCH] PatchRLStatistics

---
 unsloth/models/__init__.py |   2 +-
 unsloth/models/rl.py       | 131 +++++++++++++++++++++++++++++++++++++
 2 files changed, 132 insertions(+), 1 deletion(-)

diff --git a/unsloth/models/__init__.py b/unsloth/models/__init__.py
index 3478dfc31a..279080173f 100644
--- a/unsloth/models/__init__.py
+++ b/unsloth/models/__init__.py
@@ -20,4 +20,4 @@ from .mistral import FastMistralModel
 from .qwen2   import FastQwen2Model
 from .dpo     import PatchDPOTrainer, PatchKTOTrainer
 from ._utils  import is_bfloat16_supported
-from .rl      import PatchRL
+from .rl      import PatchRL, PatchRLStatistics
diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index b041277e47..f8d4d5412c 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -14,9 +14,22 @@
 
 __all__ = [
     "PatchRL",
+    "PatchRLStatistics",
 ]
 
 import torch
+try:
+    from transformers.utils.notebook import (
+        IntervalStrategy,
+        NotebookTrainingTracker,
+        NotebookProgressCallback,
+    )
+    HAS_NOTEBOOK = True
+except:
+    HAS_NOTEBOOK = False
+pass
+from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
+
 
 def PatchRL(FastLanguageModel):
 
@@ -43,3 +56,121 @@ def PatchRL(FastLanguageModel):
             exec(f"trl.trainer.{trainer}.{unwrap} = unsloth_{unwrap}")
     pass
 pass
+
+
+def NotebookProgressCallback_on_train_begin(Trainer_metrics):
+    def _NotebookProgressCallback_on_train_begin(self, args, state, control, **kwargs):
+        self.first_column = "Epoch" if args.eval_strategy == IntervalStrategy.EPOCH else "Step"
+        self.training_loss = 0
+        self.last_log = 0
+        column_names = [self.first_column] + ["Training Loss"]
+        if args.eval_strategy != IntervalStrategy.NO:
+            column_names.append("Validation Loss")
+        column_names += [x.replace("/", " / ") for x in Trainer_metrics]
+        self.training_tracker = NotebookTrainingTracker(state.max_steps, column_names)
+    pass
+    return _NotebookProgressCallback_on_train_begin
+pass
+
+
+def NotebookProgressCallback_on_log(Trainer_metrics):
+    def _NotebookProgressCallback_on_log(self, args, state, control, logs=None, **kwargs):
+        # Only for when there is no evaluation
+        if args.eval_strategy == IntervalStrategy.NO and "loss" in logs:
+            values = {"Training Loss": logs["loss"]}
+            for metric in DPOTrainer_metrics:
+                values[metric.replace("/", " / ")] = logs[metric]
+            pass
+            # First column is necessarily Step since we're not in epoch eval strategy
+            values["Step"] = state.global_step
+            self.training_tracker.write_line(values)
+        pass
+    pass
+    return _NotebookProgressCallback_on_log
+pass
+
+
+def _NotebookTrainingTracker_write_line(Trainer_metrics):
+    set_Trainer_metrics = set(Trainer_metrics)
+    def NotebookTrainingTracker_write_line(self, values):
+        """
+        Write the values in the inner table.
+
+        Args:
+            values (`Dict[str, float]`): The values to display.
+        """
+        if self.inner_table is None:
+            self.inner_table = [list(values.keys()), list(values.values())]
+        else:
+            columns = self.inner_table[0]
+            new_values = {}
+            for key, value in values.items():
+                lowered = key.lower()
+                if lowered in set_Trainer_metrics:
+                    new_values[lowered.replace("/", " / ")] = value
+                else:
+                    new_values[key] = value
+            pass
+            values = new_values
+
+            self.inner_table[0] = columns
+            if len(self.inner_table) > 1:
+                last_values = self.inner_table[-1]
+                first_column = self.inner_table[0][0]
+                if last_values[0] != values[first_column]:
+                    # write new line
+                    self.inner_table.append([values[c] if c in values else "No Log" for c in columns])
+                else:
+                    # update last line
+                    new_values = values
+                    for c in columns:
+                        if c not in new_values.keys():
+                            new_values[c] = last_values[columns.index(c)]
+                    self.inner_table[-1] = [new_values[c] for c in columns]
+            else:
+                # Edit for evaluation purposes
+                self.inner_table.append([values[c] if c in values else 0 for c in columns])
+            pass
+        pass
+    pass
+    return NotebookTrainingTracker_write_line
+pass
+
+
+def _PatchRLStatistics(metrics):
+    if HAS_NOTEBOOK:
+        from transformers.trainer import is_in_notebook
+        if is_in_notebook():
+            # Patch DPO notebook printing
+            NotebookTrainingTracker.write_line = NotebookTrainingTracker_write_line(metrics)
+            from transformers.trainer import DEFAULT_PROGRESS_CALLBACK
+            DEFAULT_PROGRESS_CALLBACK.on_train_begin = NotebookProgressCallback_on_train_begin(metrics)
+            DEFAULT_PROGRESS_CALLBACK.on_log         = NotebookProgressCallback_on_log(metrics)
+        pass
+    pass
+pass
+
+
+def PatchRLStatistics(algorithm = "grpo"):
+    if algorithm == "grpo":
+        metrics = [
+            "completion_length",
+            "reward",
+            "reward_std",
+            "kl",
+        ]
+    elif algorithm == "dpo" or algorithm == "kto":
+        metrics = [
+            "rewards/chosen",
+            "rewards/rejected",
+            "rewards/accuracies",
+            "rewards/margins",
+            "logps/rejected",
+            "logps/chosen",
+            "logits/rejected",
+            "logits/chosen",
+        ]
+    else:
+        print(f"Unsloth for {algorithm.upper()} is not yet implemented! Just ignore this function.")
+    _PatchRLStatistics(metrics)
+pass