Promote PPOv2Trainer and PPOv2Config to top-level import

huggingface · Sep 4, 2024 · 96ae02a · 96ae02a
1 parent fc20db8
commit 96ae02a
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 4 deletions.
diff --git a/examples/scripts/ppo/ppo.py b/examples/scripts/ppo/ppo.py
@@ -9,8 +9,7 @@
     HfArgumentParser,
 )
 
-from trl import ModelConfig
-from trl.trainer.ppov2_trainer import PPOv2Config, PPOv2Trainer
+from trl import ModelConfig, PPOv2Config, PPOv2Trainer
 from trl.trainer.utils import SIMPLE_QUERY_CHAT_TEMPLATE
 
 

diff --git a/examples/scripts/ppo/ppo_tldr.py b/examples/scripts/ppo/ppo_tldr.py
@@ -9,8 +9,7 @@
     HfArgumentParser,
 )
 
-from trl import ModelConfig
-from trl.trainer.ppov2_trainer import PPOv2Config, PPOv2Trainer
+from trl import ModelConfig, PPOv2Config, PPOv2Trainer
 from trl.trainer.utils import SIMPLE_QUERY_CHAT_TEMPLATE
 
 

diff --git a/trl/__init__.py b/trl/__init__.py
@@ -56,6 +56,8 @@
         "ORPOTrainer",
         "PPOConfig",
         "PPOTrainer",
+        "PPOv2Config",
+        "PPOv2Trainer",
         "RewardConfig",
         "RewardTrainer",
         "SFTConfig",
@@ -144,6 +146,8 @@
         ORPOTrainer,
         PPOConfig,
         PPOTrainer,
+        PPOv2Config,
+        PPOv2Trainer,
         RewardConfig,
         RewardTrainer,
         SFTConfig,

diff --git a/trl/trainer/__init__.py b/trl/trainer/__init__.py
@@ -47,6 +47,8 @@
     "orpo_trainer": ["ORPOTrainer"],
     "ppo_config": ["PPOConfig"],
     "ppo_trainer": ["PPOTrainer"],
+    "ppov2_config": ["PPOv2Config"],
+    "ppov2_trainer": ["PPOv2Trainer"],
     "reward_config": ["RewardConfig"],
     "reward_trainer": ["RewardTrainer", "compute_accuracy"],
     "sft_config": ["SFTConfig"],
@@ -112,6 +114,8 @@
     from .orpo_trainer import ORPOTrainer
     from .ppo_config import PPOConfig
     from .ppo_trainer import PPOTrainer
+    from .ppov2_config import PPOv2Config
+    from .ppov2_trainer import PPOv2Trainer
     from .reward_config import RewardConfig
     from .reward_trainer import RewardTrainer, compute_accuracy
     from .sft_config import SFTConfig