1. Introduction

In this example, we show how to use scripts to make your fine-tuning process more convenient

2. Installation

git clone https://github.com/FlagOpen/FlagEmbedding.git
cd FlagEmbedding/scripts

3. Usage

Hard Negatives

Hard negatives is a widely used method to improve the quality of sentence embedding. You can mine hard negatives following this command:

python hn_mine.py \
--input_file toy_finetune_data.jsonl \
--output_file toy_finetune_data_minedHN.jsonl \
--range_for_sampling 2-200 \
--negative_number 15 \
--use_gpu_for_searching \
--embedder_name_or_path BAAI/bge-base-en-v1.5

input_file: json data for finetuning. This script will retrieve top-k documents for each query, and random sample negatives from the top-k documents (not including the positive documents).
output_file: path to save JSON data with mined hard negatives for finetuning
negative_number: the number of sampled negatives
range_for_sampling: where to sample negative. For example, 2-100 means sampling negative_number negatives from top2-top200 documents. You can set larger value to reduce the difficulty of negatives (e.g., set it 60-300 to sample negatives from top60-300 passages)
candidate_pool: The pool to retrieval. The default value is None, and this script will retrieve from the combination of all neg in input_file. If provided, it should be a jsonl file, each line is a dict with a key text. If input a candidate_pool, this script will retrieve negatives from this file.
use_gpu_for_searching: whether to use faiss-gpu to retrieve negatives.
search_batch_size: batch size for searching. Default is 64.
embedder_name_or_path: The name or path to the embedder.
embedder_model_class: Class of the model used for embedding (current options include 'encoder-only-base', 'encoder-only-m3', 'decoder-only-base', 'decoder-only-icl'.). Default is None. For the custom model, you should set this argument.
normalize_embeddings: Set to True to normalize embeddings.
pooling_method: The pooling method for the embedder.
use_fp16: Use FP16 precision for inference.
devices: List of devices used for inference.
query_instruction_for_retrieval, query_instruction_format_for_retrieval: Instructions and format for query during retrieval.
examples_for_task, examples_instruction_format: Example tasks and their instructions format. This is only used when embedder_model_class is set to decoder-only-icl.
trust_remote_code: Set to True to trust remote code execution.
cache_dir: Cache directory for models.
embedder_batch_size: Batch sizes for embedding and reranking.
embedder_query_max_length, embedder_passage_max_length: Maximum length for embedding queries and passages.

Teacher Scores

Teacher scores can be used for model distillation. You can obtain the scores using the following command:

python add_reranker_score.py \
--input_file toy_finetune_data_minedHN.jsonl \
--output_file toy_finetune_data_score.jsonl \
--reranker_name_or_path BAAI/bge-reranker-v2-m3

input_file: path to save JSON data with mined hard negatives for finetuning
output_file: path to save JSON data with scores for finetuning
use_fp16: Whether to use fp16 for inference. Default: True
devices: Devices to use for inference. Default: None, multiple values allowed
trust_remote_code: Trust remote code. Default: False
reranker_name_or_path: The reranker name or path. Default: None
reranker_model_class: The reranker model class. Available classes: ['auto', 'encoder-only-base', 'decoder-only-base', 'decoder-only-layerwise', 'decoder-only-lightweight']. Default: auto
reranker_peft_path: The reranker peft path. Default: None
use_bf16: Whether to use bf16 for inference. Default: False
query_instruction_for_rerank: Instruction for query. Default: None
query_instruction_format_for_rerank: Format for query instruction. Default: {{}{}}
passage_instruction_for_rerank: Instruction for passage. Default: None
passage_instruction_format_for_rerank: Format for passage instruction. Default: {{}{}}
cache_dir: Cache directory for models. Default: None
reranker_batch_size: Batch size for inference. Default: 3000
reranker_query_max_length: Max length for reranking queries. Default: None
reranker_max_length: Max length for reranking. Default: 512
normalize: Whether to normalize the reranking scores. Default: False
prompt: The prompt for the reranker. Default: None
cutoff_layers: The output layers of layerwise/lightweight reranker. Default: None
compress_ratio: The compress ratio of lightweight reranker. Default: 1
compress_layers: The compress layers of lightweight reranker. Default: None, multiple values allowed

Split Data by Length

You can split the data using the following command:

python split_data_by_length.py \
--input_path train_data \
--output_dir train_data_split \
--cache_dir .cache \
--log_name .split_log \
--length_list 0 500 1000 2000 3000 4000 5000 6000 7000 \
--model_name_or_path BAAI/bge-m3 \
--num_proc 16

input_path: The path of input data. It can be a file or a directory containing multiple files.
output_dir: The directory of output data. The split data files will be saved to this directory.
cache_dir: The cache directory. Default: None
log_name: The name of the log file. Default: .split_log, which will be saved to output_dir
length_list: The length list to split. Default: [0, 500, 1000, 2000, 3000, 4000, 5000, 6000, 7000]
model_name_or_path: The model name or path of the tokenizer. Default: BAAI/bge-m3
num_proc: The number of processes. Default: 16
overwrite: Whether to overwrite the output file. Default: False