290 lines
13 KiB
Markdown
290 lines
13 KiB
Markdown
# Finetune
|
|
|
|
In this example, we show how to finetune the reranker with your data.
|
|
|
|
- [1. Installation](#1-Installation)
|
|
- [2. Data format](#2-Data-format)
|
|
- [Hard Negatives](#Hard-Negatives)
|
|
- [Teacher Scores](#Teacher-Scores)
|
|
- [3. Train](#3-Train)
|
|
- [(1) standard model](#1-standard-model)
|
|
- [(2) bge-reranker-v2-gemma](#2-bge-reranker-v2-gemma)
|
|
- [(3) bge-reranker-v2-layerwise-minicpm](#3-bge-reranker-v2-layerwise-minicpm)
|
|
|
|
## 1. Installation
|
|
|
|
- **with pip**
|
|
|
|
```shell
|
|
pip install -U FlagEmbedding[finetune]
|
|
```
|
|
|
|
- **from source**
|
|
|
|
```shell
|
|
git clone https://github.com/FlagOpen/FlagEmbedding.git
|
|
cd FlagEmbedding
|
|
pip install .[finetune]
|
|
```
|
|
|
|
For development, install as editable:
|
|
|
|
```shell
|
|
pip install -e .[finetune]
|
|
```
|
|
|
|
## 2. Data format
|
|
|
|
Train data should be a json file, where each line is a dict like this:
|
|
|
|
```shell
|
|
{"query": str, "pos": List[str], "neg":List[str], "pos_scores": List[int], "neg_scores": List[int], "prompt": str}
|
|
```
|
|
|
|
`query` is the query, and `pos` is a list of positive texts, `neg` is a list of negative texts. `pos_scores` is a list of scores corresponding to the `query` and `pos`, `neg_scores` is a list of scores corresponding to the `query` and `neg`, if you don't use knowledge distillation, it can be ignored. `prompt` is the prompt used for the input, input has the following format: `query [sep] passage [sep] prompt`. If you have no negative texts for a query, you can random sample some from the entire corpus as the negatives.
|
|
|
|
See [example_data](https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/finetune/embedder/example_data) for more detailed files.
|
|
|
|
### Hard Negatives
|
|
|
|
Hard negatives is a widely used method to improve the quality of sentence embedding. You can mine hard negatives following this command:
|
|
|
|
```shell
|
|
git clone https://github.com/FlagOpen/FlagEmbedding.git
|
|
cd FlagEmbedding/scripts
|
|
```
|
|
|
|
```shell
|
|
python hn_mine.py \
|
|
--model_name_or_path BAAI/bge-base-en-v1.5 \
|
|
--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
|
|
```
|
|
|
|
- **`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`. The format of this file is the same as [pretrain data](https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/pretrain#2-data-format). 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.
|
|
|
|
### Teacher Scores
|
|
|
|
Teacher scores can be used for model distillation. You can obtain the scores using the following command:
|
|
|
|
```shell
|
|
git clone https://github.com/FlagOpen/FlagEmbedding.git
|
|
cd FlagEmbedding/scripts
|
|
```
|
|
|
|
```shell
|
|
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 \
|
|
--devices cuda:0 cuda:1 \
|
|
--cache_dir ./cache/model \
|
|
--reranker_query_max_length 512 \
|
|
--reranker_max_length 1024
|
|
```
|
|
|
|
- **`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
|
|
|
|
## 3. Train
|
|
|
|
Detailed examples of various fine-tuning can be found in the bash files located in the corresponding folders. Here, we simply provide the training methods for the `standard model`, `bge-reranker-v2-gemma` and `bge-reranker-v2-layerwise-minicpm`.
|
|
|
|
Here are some import arguments:
|
|
|
|
- **`model_name_or_path`**: The model checkpoint for initialization.
|
|
- **`config_name`**: Pretrained config name or path if not the same as model_name. Default: None
|
|
- **`tokenizer_name`**: Pretrained tokenizer name or path if not the same as model_name. Default: None
|
|
- **`cache_dir`**: Where do you want to store the pre-trained models downloaded from s3. Default: None
|
|
- **`trust_remote_code`**: Trust remote code. Default: False
|
|
- **`model_type`**: Type of finetune, ['encoder', 'decoder']. Default: 'encoder'
|
|
- **`token`**: The token to use when accessing the model. Default: Value from environment variable HF_TOKEN or None if not set
|
|
- **`train_data`**: One or more paths to training data. `query: str`, `pos: List[str]`, `neg: List[str]` are required in the training data. Default: None
|
|
- **`cache_path`**: Where do you want to store the cached data. Default: None
|
|
- **`train_group_size`**: Default: 8
|
|
- **`query_max_len`**: The maximum total input sequence length after tokenization for passage. Sequences longer than this will be truncated. Default: 32
|
|
- **`passage_max_len`**: The maximum total input sequence length after tokenization for passage. Sequences longer than this will be truncated. Default: 128
|
|
- **`max_len`**: The maximum total input sequence length after tokenization. Sequences longer than this will be truncated. Default: 512
|
|
- **`pad_to_multiple_of`**: If set, will pad the sequence to be a multiple of the provided value. Default: None
|
|
- **`max_example_num_per_dataset`**: The max number of examples for each dataset. Default: 100000000
|
|
- **`query_instruction_for_rerank`**: Instruction for query. Default: None
|
|
- **`query_instruction_format`**: Format for query instruction. Default: "{}{}"
|
|
- **`knowledge_distillation`**: Use knowledge distillation when `pos_scores: List[float]` and `neg_scores: List[float]` are in features of training data. Default: False
|
|
- **`passage_instruction_for_rerank`**: Instruction for passage. Default: None
|
|
- **`passage_instruction_format`**: Format for passage instruction. Default: "{}{}"
|
|
- **`shuffle_ratio`**: The ratio of shuffling the text. Default: 0.0
|
|
- **`sep_token`**: The separator token for LLM reranker to discriminate between query and passage. Default: '\n'
|
|
|
|
### (1) standard model
|
|
|
|
```shell
|
|
torchrun --nproc_per_node 2 \
|
|
-m FlagEmbedding.finetune.reranker.encoder_only.base \
|
|
--model_name_or_path BAAI/bge-reranker-v2-m3 \
|
|
--cache_dir ./cache/model \
|
|
--train_data ./example_data/normal/examples.jsonl \
|
|
--cache_path ./cache/data \
|
|
--train_group_size 8 \
|
|
--query_max_len 512 \
|
|
--passage_max_len 512 \
|
|
--pad_to_multiple_of 8 \
|
|
--knowledge_distillation False \
|
|
--output_dir ./test_encoder_only_base_bge-reranker-base \
|
|
--overwrite_output_dir \
|
|
--learning_rate 6e-5 \
|
|
--fp16 \
|
|
--num_train_epochs 2 \
|
|
--per_device_train_batch_size 2 \
|
|
--gradient_accumulation_steps 1 \
|
|
--dataloader_drop_last True \
|
|
--warmup_ratio 0.1 \
|
|
--gradient_checkpointing \
|
|
--weight_decay 0.01 \
|
|
--deepspeed ../ds_stage0.json \
|
|
--logging_steps 1 \
|
|
--save_steps 1000
|
|
```
|
|
|
|
### (2) bge-reranker-v2-gemma
|
|
|
|
```shell
|
|
torchrun --nproc_per_node 2 \
|
|
-m FlagEmbedding.finetune.reranker.decoder_only.base \
|
|
--model_name_or_path BAAI/bge-reranker-v2-gemma \
|
|
--use_lora True \
|
|
--lora_rank 32 \
|
|
--lora_alpha 64 \
|
|
--use_flash_attn True \
|
|
--target_modules q_proj k_proj v_proj o_proj \
|
|
--save_merged_lora_model True \
|
|
--model_type decoder \
|
|
--cache_dir ./cache/model \
|
|
--train_data ./example_data/prompt_based/examples.jsonl \
|
|
--cache_path ./cache/data \
|
|
--train_group_size 8 \
|
|
--query_max_len 512 \
|
|
--passage_max_len 512 \
|
|
--pad_to_multiple_of 8 \
|
|
--knowledge_distillation False \
|
|
--query_instruction_for_rerank 'A: ' \
|
|
--query_instruction_format '{}{}' \
|
|
--passage_instruction_for_rerank 'B: ' \
|
|
--passage_instruction_format '{}{}' \
|
|
--output_dir ./test_decoder_only_base_bge-reranker-v2-minicpm-layerwise \
|
|
--overwrite_output_dir \
|
|
--learning_rate 2e-4 \
|
|
--bf16 \
|
|
--num_train_epochs 1 \
|
|
--per_device_train_batch_size 2 \
|
|
--gradient_accumulation_steps 1 \
|
|
--dataloader_drop_last True \
|
|
--warmup_ratio 0.1 \
|
|
--gradient_checkpointing \
|
|
--weight_decay 0.01 \
|
|
--deepspeed ../ds_stage0.json \
|
|
--logging_steps 1 \
|
|
--save_steps 1000
|
|
```
|
|
|
|
Here are some new arguments:
|
|
|
|
- **`use_lora`**: If passed, will use LORA (low-rank parameter-efficient training) to train the model.
|
|
- **`lora_rank`**: The rank of lora.
|
|
- **`lora_alpha`**: The alpha parameter of lora.
|
|
- **`lora_dropout`**: The dropout rate of lora modules.
|
|
- **`target_modules`**: The target modules to apply LORA.
|
|
- **`modules_to_save`**: List of modules that should be saved in the final checkpoint.
|
|
- **`use_flash_attn`**: If passed, will use flash attention to train the model.
|
|
- **`from_peft`**: (metadata not provided)
|
|
- **`raw_peft`**: (metadata not provided)
|
|
- **`save_merged_lora_model`**: If passed, will merge the lora modules and save the entire model.
|
|
|
|
### (3) bge-reranker-v2-layerwise-minicpm
|
|
|
|
```shell
|
|
torchrun --nproc_per_node 2 \
|
|
-m FlagEmbedding.finetune.reranker.decoder_only.layerwise \
|
|
--model_name_or_path BAAI/bge-reranker-v2-minicpm-layerwise \
|
|
--use_lora True \
|
|
--lora_rank 32 \
|
|
--lora_alpha 64 \
|
|
--use_flash_attn True \
|
|
--target_modules q_proj k_proj v_proj o_proj \
|
|
--save_merged_lora_model True \
|
|
--model_type decoder \
|
|
--model_type from_finetuned_model \
|
|
--start_layer 8 \
|
|
--head_multi True \
|
|
--head_type simple \
|
|
--trust_remote_code True \
|
|
--cache_dir ./cache/model \
|
|
--train_data ./example_data/prompt_based/examples.jsonl \
|
|
--cache_path ./cache/data \
|
|
--train_group_size 8 \
|
|
--query_max_len 512 \
|
|
--passage_max_len 512 \
|
|
--pad_to_multiple_of 8 \
|
|
--knowledge_distillation False \
|
|
--query_instruction_for_rerank 'A: ' \
|
|
--query_instruction_format '{}{}' \
|
|
--passage_instruction_for_rerank 'B: ' \
|
|
--passage_instruction_format '{}{}' \
|
|
--output_dir ./test_decoder_only_base_bge-reranker-v2-minicpm-layerwise \
|
|
--overwrite_output_dir \
|
|
--learning_rate 2e-4 \
|
|
--bf16 \
|
|
--num_train_epochs 1 \
|
|
--per_device_train_batch_size 2 \
|
|
--gradient_accumulation_steps 1 \
|
|
--dataloader_drop_last True \
|
|
--warmup_ratio 0.1 \
|
|
--gradient_checkpointing \
|
|
--weight_decay 0.01 \
|
|
--deepspeed ../ds_stage0.json \
|
|
--logging_steps 1 \
|
|
--save_steps 1000
|
|
```
|
|
|
|
Here are some new arguments:
|
|
|
|
- **`use_lora`**: If passed, will use LORA (low-rank parameter-efficient training) to train the model.
|
|
- **`lora_rank`**: The rank of lora.
|
|
- **`lora_alpha`**: The alpha parameter of lora.
|
|
- **`lora_dropout`**: The dropout rate of lora modules.
|
|
- **`target_modules`**: The target modules to apply LORA.
|
|
- **`modules_to_save`**: List of modules that should be saved in the final checkpoint.
|
|
- **`use_flash_attn`**: If passed, will use flash attention to train the model.
|
|
- **`save_merged_lora_model`**: If passed, will merge the lora modules and save the entire model.
|
|
- **`model_type`**: Model type context, which should be one of ['from_raw_model', 'from_finetuned_model'].
|
|
- **`start_layer`**: Specifies which layer to start to compute score.
|
|
- **`head_multi`**: Indicates whether to use one or multiple classifiers.
|
|
- **`head_type`**: The type of the classifier.
|