sft.py

#!/usr/bin/env python3
import argparse
import subprocess
import os

parser = argparse.ArgumentParser(description="Demo script for the model.")

parser.add_argument("--model", type=str)
parser.add_argument("--dataset", type=str)
parser.add_argument("--flavor", type=str)
parser.add_argument("--token", type=str)

# Parse the command line arguments
args = parser.parse_args()

# Get the config file based on the model and flavor
config_file = f"configs/{args.model}_{args.flavor}.yaml"

# Check if the config file exists
if not os.path.exists(config_file):
    raise RuntimeError(f"Training model {args.model} with flavor {args.flavor} is not supported.")

# Run another Python script
subprocess.run(["trl", "sft", "--config", config_file, "--dataset_name", args.dataset])