import argparse
import os

def get_args():
    parser = argparse.ArgumentParser(description="ChatQA-HF")

    ## model
    # parser.add_argument('--model-folder', type=str, default='/lustre/fsw/portfolios/llmservice/users/pengx/projects/vllm_run/')
    # parser.add_argument('--model-name', type=str, default='Llama-3-70B-Instruct-Gradient-262k')

    parser.add_argument('--model-folder', type=str, default='/lustre/fsw/portfolios/llmservice/users/pengx/projects/swa_long_pretrain_llama2/checkpoints/applications/long_131072_25_multiturn_qa_blend_commercial_v28_9_multiturn_pp1_hf')
    parser.add_argument('--model-name', type=str, default='ChatQA2')

    ## tokenizer
    # parser.add_argument('--tokenizer-path', type=str, default='/lustre/fsw/portfolios/adlr/users/zihanl/inform/ckpts/llama2-tokenizer')
    parser.add_argument('--tokenizer-path', type=str, default='/lustre/fsw/portfolios/llmservice/users/pengx/projects/vllm_run/Llama-3-70B-Instruct-Gradient-262k/')
    # parser.add_argument('--tokenizer-path', type=str, default='/lustre/fsw/portfolios/llmservice/users/pengx/projects/swa_long_pretrain_llama2/checkpoints/applications/long_131072_25_multiturn_qa_blend_commercial_v28_9_multiturn_pp1_hf')

    ## dataset path
    # parser.add_argument('--data-folder', type=str, default='/lustre/fsw/portfolios/adlr/users/zihanl/datasets/foundational_qa/test_benchmarks/multi-turn-qa')
    parser.add_argument('--data-folder', type=str, default='/lustre/fs1/portfolios/llmservice/users/pengx/projects/vllm_run/oss_test/')
    parser.add_argument('--data-folder-singleturn', type=str, default='/lustre/fsw/portfolios/adlr/users/zihanl/datasets/foundational_qa/test_benchmarks/single-turn-qa')
    parser.add_argument('--data-folder-scrolleval', type=str, default='/lustre/fsw/portfolios/adlr/users/zihanl/datasets/foundational_qa/scroll_eval_data')

    parser.add_argument('--eval-dataset', type=str, default='')
    # parser.add_argument('--doc2dial-path', type=str, default='doc2dial/doc2dial_ftdragon_chatgptgen7k_chunk150_QA_test.json')
    # parser.add_argument('--convfinqa-path', type=str, default='convfinqav3/convfinqav3_QA_dev.json')
    # parser.add_argument('--convfinqa-path', type=str, default='convfinqa_general/convfinqa_general_QA_dev.json')
    # parser.add_argument('--quac-path', type=str, default='quac/quac_ftdragon_chatgptgen7k_chunk150_QA_test.json')
    # parser.add_argument('--qrecc-path', type=str, default='qrecc/qrecc_ftdragon_chatgptgen7k_chunk150_QA_test.json')
    # parser.add_argument('--doqa-cooking-path', type=str, default='doqa/doqa_cooking_QA_test.json')
    # parser.add_argument('--doqa-travel-path', type=str, default='doqa/doqa_travel_QA_test.json')
    # parser.add_argument('--doqa-movies-path', type=str, default='doqa/doqa_movies_QA_test.json')
    # parser.add_argument('--coqa-path', type=str, default='coqa/coqa_QA_dev.json')
    # # parser.add_argument('--hybridial-path', type=str, default='HybridDial/HybridDial_fqa_test.json')
    # parser.add_argument('--hybridial-path', type=str, default='HybridDial_general/HybridDial_general_QA_test.json')
    # # parser.add_argument('--sqa-path', type=str, default='sqa/sqa_QA_test.json')
    # parser.add_argument('--sqa-path', type=str, default='sqa_general/sqa_general_QA_test.json')
    # parser.add_argument('--topiocqa-path', type=str, default='topiocqa/topiocqa_dev_retrieval_dragon_ft_chatgptgen7k.json')
    # parser.add_argument('--inscit-path', type=str, default='inscit/inscit_dev_retrieval_dragon_ft_chatgptgen7k_with_topic.json')

    parser.add_argument('--doc2dial-path', type=str, default='doc2dial/test.json')
    parser.add_argument('--convfinqa-path', type=str, default='convfinqa/dev.json')
    parser.add_argument('--quac-path', type=str, default='quac/test.json')
    parser.add_argument('--qrecc-path', type=str, default='qrecc/test.json')
    parser.add_argument('--doqa-cooking-path', type=str, default='doqa/test_cooking.json')
    parser.add_argument('--doqa-travel-path', type=str, default='doqa/test_travel.json')
    parser.add_argument('--doqa-movies-path', type=str, default='doqa/test_movies.json')
    parser.add_argument('--coqa-path', type=str, default='coqa/dev.json')
    parser.add_argument('--hybridial-path', type=str, default='hybridial/test.json')
    parser.add_argument('--sqa-path', type=str, default='sqa/test.json')
    parser.add_argument('--topiocqa-path', type=str, default='topiocqa/dev.json')
    parser.add_argument('--inscit-path', type=str, default='inscit/dev.json')
    
    parser.add_argument('--kilt-nq-path', type=str, default='kilt/nq/test.json')
    parser.add_argument('--kilt-tqa-path', type=str, default='kilt/tqa/test.json')
    parser.add_argument('--kilt-hotpotqa-path', type=str, default='kilt/hotpotqa/test.json')
    # parser.add_argument('--kilt-hotpotqa-path', type=str, default='kilt/hotpotqa_rerank/test.json')

    parser.add_argument('--nq-path', type=str, default='nq_dragon_retrieved/test.json')
    parser.add_argument('--tqa-path', type=str, default='triviaqa_dragon_retrieved/test.json')
    parser.add_argument('--hotpotqa-path', type=str, default='hotpotqa_dragon_retrieved/test.json')

    ## scroll eval
    parser.add_argument('--scroll-hotpotqa-chunk1200-path', type=str, default='hotpotqa.e5_mistral_retriever_chunkbysents1200/test.json')
    parser.add_argument('--scroll-musique-chunk1200-path', type=str, default='musique.e5_mistral_retriever_chunkbysents1200/test.json')
    parser.add_argument('--scroll-qasper-chunk1200-path', type=str, default='qasper.e5_mistral_retriever_chunkbysents1200/test.json')
    parser.add_argument('--scroll-narrative_qa-chunk1200-path', type=str, default='narrative_qa.e5_mistral_retriever_chunkbysents1200/test.json')
    parser.add_argument('--scroll-quality-chunk1200-path', type=str, default='quality.e5_mistral_retriever_chunkbysents1200/test.json')
    parser.add_argument('--scroll-multifieldqa_en-chunk1200-path', type=str, default='multifieldqa_en.e5_mistral_retriever_chunkbysents1200/test.json')
    parser.add_argument('--scroll-qmsum-chunk1200-path', type=str, default='qmsum.e5_mistral_retriever_chunkbysents1200/test.json')

    parser.add_argument('--scroll-hotpotqa-chunk300-path', type=str, default='hotpotqa.e5_mistral_retriever_chunkbysents1200/test.json')
    parser.add_argument('--scroll-musique-chunk300-path', type=str, default='musique.e5_mistral_retriever_chunkbysents1200/test.json')
    parser.add_argument('--scroll-qasper-chunk300-path', type=str, default='qasper.e5_mistral_retriever_chunkbysents1200/test.json')
    parser.add_argument('--scroll-narrative_qa-chunk300-path', type=str, default='narrative_qa.e5_mistral_retriever_chunkbysents1200/test.json')
    parser.add_argument('--scroll-quality-chunk300-path', type=str, default='quality.e5_mistral_retriever_chunkbysents1200/test.json')
    parser.add_argument('--scroll-multifieldqa_en-chunk300-path', type=str, default='multifieldqa_en.e5_mistral_retriever_chunkbysents1200/test.json')
    parser.add_argument('--scroll-qmsum-chunk300-path', type=str, default='qmsum.e5_mistral_retriever_chunkbysents1200/test.json')

    parser.add_argument('--sample-input-file', type=str, default='')
    parser.add_argument("--use-retrieved-neighbours", action='store_true', default=False,
                       help='Use retrieved neighbours')

    ## others
    parser.add_argument('--max-seq-length', type=int, default=128000)
    parser.add_argument('--num-ctx', type=int, default=5)
    parser.add_argument('--start-idx', type=int, default=-1)
    parser.add_argument('--end-idx', type=int, default=-1)
    parser.add_argument('--max-tokens', type=int, default=64)

    args = parser.parse_args()

    return args