import argparse

from colossalqa.retrieval_conversation_universal import UniversalRetrievalConversation

if __name__ == "__main__":
    # Parse arguments
    parser = argparse.ArgumentParser()
    parser.add_argument("--en_model_path", type=str, default=None)
    parser.add_argument("--zh_model_path", type=str, default=None)
    parser.add_argument("--zh_model_name", type=str, default=None)
    parser.add_argument("--en_model_name", type=str, default=None)
    parser.add_argument(
        "--sql_file_path", type=str, default=None, help="path to the a empty folder for storing sql files for indexing"
    )
    args = parser.parse_args()

    # Will ask for documents path in running time
    session = UniversalRetrievalConversation(
        files_en=None,
        files_zh=None,
        zh_model_path=args.zh_model_path,
        en_model_path=args.en_model_path,
        zh_model_name=args.zh_model_name,
        en_model_name=args.en_model_name,
        sql_file_path=args.sql_file_path,
    )
    session.start_test_session()