RL metrics

unsloth/models/dpo.py

@@ -17,8 +17,8 @@ __all__ = [
     "PatchKTOTrainer",
 ]
 
-from .rl import PatchRLStatistics
+from .rl import PatchFastRL
 
-def PatchDPOTrainer(): PatchRLStatistics("DPO")
+def PatchDPOTrainer(): PatchFastRL("DPO")
 
-def PatchKTOTrainer(): PatchRLStatistics("KTO")
+def PatchKTOTrainer(): PatchFastRL("KTO")

unsloth/models/rl.py

@@ -13,8 +13,7 @@
 # limitations under the License.
 
 __all__ = [
-    "PatchRL",
-    "PatchRLStatistics",
+    "PatchFastRL",
 ]
 
 import torch
@@ -202,6 +201,9 @@ def get_trl_metrics():
         left_prefix = 'prefix = "eval_" if train_eval == "eval" else ""' in file
         if left_prefix: metrics += metrics_f
 
+        # Remove all eval_ things
+        metrics = [x for x in metrics if not x.startswith("eval_")]
+
         all_metrics[trainer[:trainer.find("_")].upper()] = metrics
     pass
     return all_metrics
@@ -219,3 +221,9 @@ def PatchRLStatistics(algorithm = "GRPO"):
     pass
     _PatchRLStatistics(all_metrics[algorithm], algorithm)
 pass
+
+
+def PatchFastRL(algorithm = "GRPO", FastLanguageModel = None):
+    if FastLanguageModel is not None: PatchRL(FastLanguageModel)
+    PatchRLStatistics(algorithm)
+pass