File size: 771 Bytes
e43de7f
411ce45
c5d1a86
04db925
411ce45
 
 
 
 
75955e4
04db925
411ce45
04db925
411ce45
04db925
 
 
 
 
 
 
c5d1a86
 
04db925
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
#!/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])