54 lines
1.6 KiB
Python
54 lines
1.6 KiB
Python
import sys
|
|
import yaml
|
|
|
|
def main():
|
|
run_type = sys.argv[1]
|
|
model = sys.argv[2]
|
|
max_steps = sys.argv[3]
|
|
run_name = sys.argv[4]
|
|
output_dir = sys.argv[5]
|
|
|
|
if run_type == "lora_sft":
|
|
yaml_file = './results/lora_sft_template.yaml'
|
|
elif run_type == "inference":
|
|
yaml_file = './results/predict_template.yaml'
|
|
|
|
model_name_or_path = ""
|
|
template = ""
|
|
if model == "9g-8B":
|
|
model_name_or_path = "../../models/sft_8b_v2"
|
|
template = ""
|
|
elif model == "Baichuan2-7B":
|
|
model_name_or_path = "../../models/Baichuan2-7B-Base"
|
|
template = "baichuan2"
|
|
elif model == "ChatGLM2-6B":
|
|
model_name_or_path = "../../models/chatglm2-6b"
|
|
template = "chatglm2"
|
|
elif model == "Llama2-7B":
|
|
model_name_or_path = "../../models/llama-2-7b-ms"
|
|
template = "llama2"
|
|
elif model == "Qwen-7B":
|
|
model_name_or_path = "../../models/Qwen-7B"
|
|
template = "qwen"
|
|
else:
|
|
print("ERROR: model not supported.")
|
|
sys.exit()
|
|
|
|
config = None
|
|
with open(yaml_file, 'r', encoding='utf-8') as f:
|
|
config = yaml.load(f.read(), Loader=yaml.FullLoader)
|
|
|
|
|
|
config['model_name_or_path'] = model_name_or_path
|
|
config['template'] = template
|
|
config['output_dir'] = output_dir
|
|
if run_type == "lora_sft":
|
|
config['max_steps'] = max_steps
|
|
|
|
with open(f'{output_dir}/{run_name}.yaml', 'w', encoding='utf-8') as f:
|
|
yaml.dump(data=config, stream=f, allow_unicode=True)
|
|
|
|
print(f"yaml file saved to {output_dir}/{run_name}.yaml")
|
|
|
|
if __name__ == "__main__":
|
|
main() |