Package: clddp
clddp: Contrastive Learning with Distributed Data ParallelThis Python package provides an implementation for constrastive learning with multiple GPUs (i.e. Distributed-Data-Parallel), with a special focus on neural retrieval.Installationpipinstall-UclddpIf ColBERT is going to be used, please install its package additionally:pipinstallgit+https://github.com/stanford-futuredata/ColBERT.git@21b460a606bed606e8a7fa105ada36b18e8084ecQuick StartPlease have a look at theexamplesfor a quick start. For example, one can run the multi-GPU training with the following script:Training exampledataset="fiqa"exportDATA_DIR="data"mkdir-p$DATA_DIRexportDATASET_PATH="$DATA_DIR/$dataset"if[!-f"$DATASET_PATH.zip"];thenwgethttps://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/$dataset.zip-P$DATA_DIRfiif[!-d"$DATASET_PATH"];thenunzip$DATASET_PATH.zip-d$DATA_DIRfiexportCUDA_VISIBLE_DEVICES="0,1,2,3"exportWANDB_MODE="online"exportCHECKPOINT_DIR="checkpoints"exportCLI_ARGS="--project="clddp_examples"--checkpoint_dir=$CHECKPOINT_DIR--query_model_name_or_path="sentence-transformers/msmarco-distilbert-base-tas-b"--shared_encoder=True--sep=blank--pooling=cls--similarity_function=dot_product--query_max_length=350--passage_max_length=350--sim_scale=1.0--fp16=True--train_data=$DATASET_PATH--train_dataloader=beir--num_negatives=0--dev_data=$DATASET_PATH--dev_dataloader=beir--do_dev=True--quick_dev=False--test_data=$DATASET_PATH--test_dataloader=beir--num_train_epochs=1--eval_steps=0.4--save_steps=0.4"exportOUTPUT_DIR=$(python-mclddp.args.train$CLI_ARGS)mkdir-p$OUTPUT_DIRexportLOG_PATH="$OUTPUT_DIR/logging.log"echo"Logging file path:$LOG_PATH"nohuptorchrun--nproc_per_node=4--master_port=29501-mclddp.train$CLI_ARGS>$LOG_PATH&Search ExampleThis will run exact search with multiple GPUs and output the retrieval results in the TREC format (each row isquery-idQ0passage-idrankscoreexp):# Build a retriever checkpoint:exportCHECKPOINT_DIR="checkpoints/off-the-shell/tasb"exportCLI_ARGS="--output_dir=$CHECKPOINT_DIR--query_model_name_or_path="sentence-transformers/msmarco-distilbert-base-tas-b"--shared_encoder=True--sep=blank--pooling=cls--similarity_function=dot_product--query_max_length=350--passage_max_length=350--sim_scale=1.0"python-mclddp.retriever$CLI_ARGS# Run search:dataset="fiqa"exportDATA_DIR="data"exportDATASET_PATH="$DATA_DIR/$dataset"if[!-d$DATASET_PATH];thenecho"Data path$DATASET_PATHdoes not exist"exitfiif[!-d$CHECKPOINT_DIR];thenecho"Checkpoint path$CHECKPOINT_DIRdoes not exist"exitfiexportCUDA_VISIBLE_DEVICES="0,1,2,3"exportCLI_ARGS="--checkpoint_dir=$CHECKPOINT_DIR--data_dir=$DATASET_PATH--dataloader=beir"exportOUTPUT_DIR=$(python-mclddp.args.search$CLI_ARGS)mkdir-p$OUTPUT_DIRexportLOG_PATH="$OUTPUT_DIR/logging.log"echo"Logging file path:$LOG_PATH"nohuptorchrun--nproc_per_node=4--master_port=29501-mclddp.search$CLI_ARGS>$LOG_PATH&Evaluation is similar and the example can be foundhere.Custom DataFor loading custom data, one needs to inherit theclddp.dataloader.BaseDataLoaderand add it to the lookup mapclddp.DATA_LOADER_LOOKUP:fromclddp.trainimportmainfromclddp.dataloaderimportBaseDataLoader,DATA_LOADER_LOOKUPfromclddp.dmimportRetrievalDatasetclassMyDataLoader(BaseDataLoader):defload_data(data_name_or_path:str,progress_bar:bool)->RetrievalDataset:...DATA_LOADER_LOOKUP["my_dataloader"]=MyDataLoaderif__name__=="__main__":main()# Same for other entry pointsThen, one can specify thexxx_dataloader=my_dataloaderin the CLI arguments.