import os

os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'
os.environ['ASCEND_RT_VISIBLE_DEVICES'] = '0,1'


def test_sft():
    from swift.megatron import MegatronSftArguments, megatron_sft_main
    megatron_sft_main(
        MegatronSftArguments(
            mcore_model='Qwen2-7B-Instruct-mcore',
            dataset=[
                'AI-ModelScope/alpaca-gpt4-data-zh#500', 'swift/self-cognition#500',
                'AI-ModelScope/alpaca-gpt4-data-en#500'
            ],
            split_dataset_ratio=0.01,
            tensor_model_parallel_size=2,
            train_iters=100,
            model_author=['swift'],
            model_name=['swift-robot'],
            sequence_parallel=True,
            finetune=True))


def test_pt():
    from swift.megatron import MegatronPretrainArguments, megatron_pretrain_main
    megatron_pretrain_main(
        MegatronPretrainArguments(
            mcore_model='Qwen2-7B-mcore',
            dataset=['AI-ModelScope/alpaca-gpt4-data-zh#500', 'AI-ModelScope/alpaca-gpt4-data-en#500'],
            split_dataset_ratio=0.01,
            tensor_model_parallel_size=2,
            train_iters=200,
            eval_iters=5,
            finetune=True))


if __name__ == '__main__':
    test_sft()
    # test_pt()