import os

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


def test_dpo():
    from swift.megatron import MegatronRLHFArguments, megatron_rlhf_main
    megatron_rlhf_main(
        MegatronRLHFArguments(
            mcore_model='Qwen2.5-3B-Instruct-mcore',
            dataset=['hjh0119/shareAI-Llama3-DPO-zh-en-emoji#10000'],
            split_dataset_ratio=0.01,
            micro_batch_size=16,
            tensor_model_parallel_size=2,
            eval_steps=5,
            logging_steps=1,
            finetune=True,
            num_train_epochs=1,
        ))


def test_hf():
    from swift import RLHFArguments, rlhf_main
    rlhf_main(
        RLHFArguments(
            model='Qwen/Qwen2.5-3B-Instruct',
            dataset=['hjh0119/shareAI-Llama3-DPO-zh-en-emoji#1000'],
            split_dataset_ratio=0.01,
            max_steps=100,
            padding_free=True,
            attn_impl='flash_attn',
            train_dataloader_shuffle=False,
            use_logits_to_keep=False,
        ))


if __name__ == '__main__':
    test_dpo()
    # test_hf()