populate SUPPORTED_COMMANDS cli

tests/test_cli.py

@@ -14,6 +14,10 @@
 import subprocess
 import sys
 import unittest
+from pathlib import Path
+import os
+import glob
+from trl.commands.cli_utils import populate_supported_commands
 
 
 @unittest.skipIf(sys.platform.startswith("win"), "Skipping on Windows")
@@ -43,3 +47,22 @@ def test_dpo_cli():
 def test_env_cli():
  output = subprocess.run("trl env", capture_output=True, text=True, shell=True, check=True)
  assert "- Python version: " in output.stdout
+
+
+def test_populate_supported_commands():
+ commands = populate_supported_commands()
+
+ # Check for specific commands
+ assert 'sft' in commands, "SFT command not found"
+ assert 'dpo' in commands, "DPO command not found"
+
+ # Check that all commands are strings and don't have .py extension
+ for cmd in commands:
+ assert isinstance(cmd, str), f"Command {cmd} is not a string"
+ assert not cmd.endswith('.py'), f"Command {cmd} should not have .py extension"
+
+ # Check that the number of commands matches the number of .py files in the scripts directory
+ trl_dir = Path(__file__).resolve().parent.parent
+ scripts_path = os.path.join(trl_dir, 'examples', 'scripts', '*.py')
+ py_files = glob.glob(scripts_path)
+ assert len(commands) == len(py_files), f"Number of commands ({len(commands)}) doesn't match number of .py files ({len(py_files)})"

trl/commands/cli.py

@@ -31,11 +31,10 @@
  is_liger_kernel_available,
  is_llmblender_available,
 )
-from .cli_utils import get_git_commit_hash
+from .cli_utils import get_git_commit_hash, populate_supported_commands
 
 
-SUPPORTED_COMMANDS = ["sft", "dpo", "chat", "kto", "env"]
-
+SUPPORTED_COMMANDS = populate_supported_commands()
 
 def print_env():
  accelerate_config = accelerate_config_str = "not found"

trl/commands/cli_utils.py

@@ -17,8 +17,10 @@
 import inspect
 import logging
 import os
+import glob
 import subprocess
 import sys
+from pathlib import Path
 from argparse import Namespace
 from dataclasses import dataclass, field
 
@@ -305,3 +307,16 @@ def get_git_commit_hash(package_name):
  return None
  except Exception as e:
  return f"Error: {str(e)}"
+
+
+def populate_supported_commands():
+ # Path to the script examples directory
+ trl_dir = Path(__file__).resolve().parent.parent.parent
+ scripts_path = os.path.join(trl_dir, 'examples', 'scripts', '*.py')
+ # find all the scripts in the examples directory
+ trainer_files = glob.glob(scripts_path)
+
+ # Extract command names without the .py extension
+ commands = [os.path.basename(f).replace('.py', '') for f in trainer_files]
+
+ return commands