chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,218 @@
|
||||
# Finetune
|
||||
In this example, we show how to finetune the baai-general-embedding with your data.
|
||||
|
||||
## 1. Installation
|
||||
```
|
||||
git clone https://github.com/FlagOpen/FlagEmbedding.git
|
||||
cd FlagEmbedding/research/baai_general_embedding
|
||||
```
|
||||
## 2. Data format
|
||||
Train data should be a json file, where each line is a dict like this:
|
||||
|
||||
```
|
||||
{"query": str, "pos": List[str], "neg":List[str]}
|
||||
```
|
||||
|
||||
`query` is the query, and `pos` is a list of positive texts, `neg` is a list of negative texts.
|
||||
If you have no negative texts for a query, you can random sample some from the entire corpus as the negatives.
|
||||
|
||||
See [toy_finetune_data.jsonl](https://github.com/FlagOpen/FlagEmbedding/blob/master/examples/finetune/toy_finetune_data.jsonl) for a toy data file.
|
||||
|
||||
### 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
|
||||
```
|
||||
|
||||
```bash
|
||||
python -m FlagEmbedding.baai_general_embedding.finetune.hn_mine \
|
||||
--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.
|
||||
|
||||
|
||||
## 3. Train
|
||||
```
|
||||
torchrun --nproc_per_node {number of gpus} \
|
||||
-m finetune.run \
|
||||
--output_dir {path to save model} \
|
||||
--model_name_or_path BAAI/bge-large-zh-v1.5 \
|
||||
--train_data ./toy_finetune_data.jsonl \
|
||||
--learning_rate 1e-5 \
|
||||
--fp16 \
|
||||
--num_train_epochs 5 \
|
||||
--per_device_train_batch_size {large batch size; set 1 for toy data} \
|
||||
--dataloader_drop_last True \
|
||||
--normlized True \
|
||||
--temperature 0.02 \
|
||||
--query_max_len 64 \
|
||||
--passage_max_len 256 \
|
||||
--train_group_size 2 \
|
||||
--negatives_cross_device \
|
||||
--logging_steps 10 \
|
||||
--save_steps 1000 \
|
||||
--query_instruction_for_retrieval ""
|
||||
```
|
||||
|
||||
**some important arguments**:
|
||||
- `per_device_train_batch_size`: batch size in training. In most of cases, larger batch size will bring stronger performance. You can expand it by enabling `--fp16`, `--deepspeed ./df_config.json` (df_config.json can refer to [ds_config.json](./ds_config.json)), `--gradient_checkpointing`, etc.
|
||||
- `train_group_size`: the number of positive and negatives for a query in training.
|
||||
There are always one positive, so this argument will control the number of negatives (#negatives=train_group_size-1).
|
||||
Noted that the number of negatives should not be larger than the numbers of negatives in data `"neg":List[str]`.
|
||||
Besides the negatives in this group, the in-batch negatives also will be used in fine-tuning.
|
||||
- `negatives_cross_device`: share the negatives across all GPUs. This argument will extend the number of negatives.
|
||||
- `learning_rate`: select a appropriate for your model. Recommend 1e-5/2e-5/3e-5 for large/base/small-scale.
|
||||
- `temperature`: It will influence the distribution of similarity scores. **Recommended value: 0.01-0.1.**
|
||||
- `query_max_len`: max length for query. Please set it according the average length of queries in your data.
|
||||
- `passage_max_len`: max length for passage. Please set it according the average length of passages in your data.
|
||||
- `query_instruction_for_retrieval`: instruction for query, which will be added to each query. You also can set it `""` to add nothing to query.
|
||||
- `use_inbatch_neg`: use passages in the same batch as negatives. Default value is True.
|
||||
- `save_steps`: for setting how many training steps to save a checkpoint.
|
||||
|
||||
For more training arguments please refer to [transformers.TrainingArguments](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments)
|
||||
|
||||
|
||||
### 4. Model merging via [LM-Cocktail](https://github.com/FlagOpen/FlagEmbedding/tree/master/research/LM_Cocktail) [optional]
|
||||
|
||||
For more details please refer to [LM-Cocktail](https://github.com/FlagOpen/FlagEmbedding/tree/master/research/LM_Cocktail).
|
||||
|
||||
Fine-tuning the base bge model can improve its performance on target task,
|
||||
but maybe lead to severe degeneration of model’s general capabilities
|
||||
beyond the targeted domain (e.g., lower performance on c-mteb tasks).
|
||||
By merging the fine-tuned model and the base model,
|
||||
LM-Cocktail can significantly enhance performance in downstream task
|
||||
while maintaining performance in other unrelated tasks.
|
||||
|
||||
```python
|
||||
from LM_Cocktail import mix_models, mix_models_with_data
|
||||
|
||||
# Mix fine-tuned model and base model; then save it to output_path: ./mixed_model_1
|
||||
model = mix_models(
|
||||
model_names_or_paths=["BAAI/bge-large-en-v1.5", "your_fine-tuned_model"],
|
||||
model_type='encoder',
|
||||
weights=[0.5, 0.5], # you can change the weights to get a better trade-off.
|
||||
output_path='./mixed_model_1')
|
||||
```
|
||||
|
||||
If you have a new task, and there is no data or resource can be used for fine-tuning,
|
||||
you can try to use LM-Cocktail to merge existing models (from open-source community or your models fine-tuned on other tasks) to produce a task-specific model.
|
||||
In this way, you just need to construct a few example data and don't need fine-tuning the base model.
|
||||
For example, you can merge the models from [huggingface](https://huggingface.co/Shitao) using the example data for your task:
|
||||
```python
|
||||
from LM_Cocktail import mix_models, mix_models_with_data
|
||||
|
||||
example_data = [
|
||||
{"query": "How does one become an actor in the Telugu Film Industry?", "pos": [" How do I become an actor in Telugu film industry?"], "neg": [" What is the story of Moses and Ramesses?", " Does caste system affect economic growth of India?"]},
|
||||
{"query": "Why do some computer programmers develop amazing software or new concepts, while some are stuck with basic programming work?", "pos": [" Why do some computer programmers develops amazing softwares or new concepts, while some are stuck with basics programming works?"], "neg": [" When visiting a friend, do you ever think about what would happen if you did something wildly inappropriate like punch them or destroy their furniture?", " What is the difference between a compliment and flirting?"]}
|
||||
]
|
||||
|
||||
model = mix_models_with_data(
|
||||
model_names_or_paths=["BAAI/bge-base-en-v1.5", "Shitao/bge-hotpotqa", "Shitao/bge-quora"],
|
||||
model_type='encoder',
|
||||
example_ata=example_data,
|
||||
temperature=5.0,
|
||||
max_input_length=512,
|
||||
neg_number=2)
|
||||
```
|
||||
**Since there are only 9 `bge-*` models in this [repo](https://huggingface.co/Shitao), the performance may not be satisfactory when your task is different with all 9 fine-tuning tasks.
|
||||
You can fine-tune the base model on more tasks and merge them to achieve better performance on your task.**
|
||||
|
||||
|
||||
### 5. Load your model
|
||||
After fine-tuning BGE model, you can load it easily in the same way as [here](https://github.com/FlagOpen/FlagEmbedding/tree/master/research/baai_general_embedding#usage)
|
||||
|
||||
Please replace the `query_instruction_for_retrieval` with your instruction if you set a different value for hyper-parameter `--query_instruction_for_retrieval` when fine-tuning.
|
||||
|
||||
|
||||
### 6. Evaluate model
|
||||
We provide [a simple script](https://github.com/FlagOpen/FlagEmbedding/blob/master/research/baai_general_embedding/finetune/eval_msmarco.py) to evaluate the model's performance.
|
||||
A brief summary of how the script works:
|
||||
|
||||
1. Load the model on all available GPUs through [DataParallel](https://pytorch.org/docs/stable/generated/torch.nn.DataParallel.html).
|
||||
2. Encode the corpus and offload the embeddings in `faiss` Flat index. By default, `faiss` also dumps the index on all available GPUs.
|
||||
3. Encode the queries and search `100` nearest neighbors for each query.
|
||||
4. Compute Recall and MRR metrics.
|
||||
|
||||
First, install `faiss`, a popular approximate nearest neighbor search library:
|
||||
```bash
|
||||
conda install -c conda-forge faiss-gpu
|
||||
```
|
||||
|
||||
#### 6.1 MSMARCO dataset
|
||||
The default evaluate data is MSMARCO, a widely used retrieval benchmark.
|
||||
|
||||
You can check the data formats for the [msmarco corpus](https://huggingface.co/datasets/namespace-Pt/msmarco-corpus) and [evaluation queries](https://huggingface.co/datasets/namespace-Pt/msmarco).
|
||||
|
||||
Run the following command:
|
||||
|
||||
```bash
|
||||
python -m finetune.eval_msmarco \
|
||||
--encoder BAAI/bge-base-en-v1.5 \
|
||||
--fp16 \
|
||||
--add_instruction \
|
||||
--k 100
|
||||
```
|
||||
**some important arguments:**
|
||||
- `encoder`: specify the encoder model, which can be either a model on huggingface or a local one.
|
||||
- `fp16`: use half precision for inference.
|
||||
- `add_instruction`: add retrieval instruction (`Represent this sentence for searching relevant passages: `).
|
||||
- `k`: specify how many nearest neighbors to retrieve for each query.
|
||||
|
||||
The results should be similar to
|
||||
```python
|
||||
{
|
||||
'MRR@1': 0.2330945558739255,
|
||||
'MRR@10': 0.35786976395142633,
|
||||
'MRR@100': 0.3692618036917553,
|
||||
'Recall@1': 0.22606255969436478,
|
||||
'Recall@10': 0.6412965616045848,
|
||||
'Recall@100': 0.9012774594078318
|
||||
}
|
||||
```
|
||||
|
||||
#### 6.2 Your dataset
|
||||
|
||||
You should prepare two files with jsonl format:
|
||||
- One is corpus_data, which contains the text you want to search. A toy example: [toy_corpus.json](./toy_evaluation_data/toy_corpus.json)
|
||||
```
|
||||
{"content": "A is ..."}
|
||||
{"content": "B is ..."}
|
||||
{"content": "C is ..."}
|
||||
{"content": "Panda is ..."}
|
||||
{"content": "... is A"}
|
||||
```
|
||||
- The other is query_data, which contains the queries and the ground truth. A toy example: [toy_corpus.json](./toy_evaluation_data/toy_query.json)
|
||||
```
|
||||
{"query": "What is A?", "positive": ["A is ...", "... is A"]}
|
||||
{"query": "What is B?", "positive": ["B is ..."]}
|
||||
{"query": "What is C?", "positive": ["C is ..."]}
|
||||
```
|
||||
|
||||
Then, pass the data path to evaluation script:
|
||||
```bash
|
||||
python -m FlagEmbedding.baai_general_embedding.finetune.eval_msmarco \
|
||||
--encoder BAAI/bge-base-en-v1.5 \
|
||||
--fp16 \
|
||||
--add_instruction \
|
||||
--k 100 \
|
||||
--corpus_data ./toy_evaluation_data/toy_corpus.json \
|
||||
--query_data ./toy_evaluation_data/toy_query.json
|
||||
```
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
{
|
||||
"fp16": {
|
||||
"enabled": "auto",
|
||||
"loss_scale": 0,
|
||||
"loss_scale_window": 1000,
|
||||
"initial_scale_power": 12,
|
||||
"hysteresis": 2,
|
||||
"min_loss_scale": 1
|
||||
},
|
||||
|
||||
"bf16": {
|
||||
"enabled": "auto"
|
||||
},
|
||||
|
||||
"optimizer": {
|
||||
"type": "AdamW",
|
||||
"params": {
|
||||
"lr": "auto",
|
||||
"betas": "auto",
|
||||
"eps": "auto",
|
||||
"weight_decay": "auto"
|
||||
}
|
||||
},
|
||||
|
||||
"scheduler": {
|
||||
"type": "WarmupDecayLR",
|
||||
"params": {
|
||||
"warmup_min_lr": "auto",
|
||||
"warmup_max_lr": "auto",
|
||||
"warmup_num_steps": "auto",
|
||||
"total_num_steps": "auto"
|
||||
}
|
||||
},
|
||||
|
||||
"zero_optimization": {
|
||||
"stage": 0
|
||||
},
|
||||
|
||||
"gradient_accumulation_steps": "auto",
|
||||
"gradient_clipping": "auto",
|
||||
"steps_per_print": 100,
|
||||
"train_batch_size": "auto",
|
||||
"train_micro_batch_size_per_gpu": "auto",
|
||||
"wall_clock_breakdown": false
|
||||
}
|
||||
@@ -0,0 +1,5 @@
|
||||
{"content": "A is ..."}
|
||||
{"content": "B is ..."}
|
||||
{"content": "C is ..."}
|
||||
{"content": "Panda is ..."}
|
||||
{"content": "... is A"}
|
||||
@@ -0,0 +1,3 @@
|
||||
{"query": "What is A?", "positive": ["A is ...", "... is A"]}
|
||||
{"query": "What is B?", "positive": ["B is ..."]}
|
||||
{"query": "What is C?", "positive": ["C is ..."]}
|
||||
@@ -0,0 +1,10 @@
|
||||
{"query": "Five women walk along a beach wearing flip-flops.", "pos": ["Some women with flip-flops on, are walking along the beach"], "neg": ["The 4 women are sitting on the beach.", "There was a reform in 1996.", "She's not going to court to clear her record.", "The man is talking about hawaii.", "A woman is standing outside.", "The battle was over. ", "A group of people plays volleyball."]}
|
||||
{"query": "A woman standing on a high cliff on one leg looking over a river.", "pos": ["A woman is standing on a cliff."], "neg": ["A woman sits on a chair.", "George Bush told the Republicans there was no way he would let them even consider this foolish idea, against his top advisors advice.", "The family was falling apart.", "no one showed up to the meeting", "A boy is sitting outside playing in the sand.", "Ended as soon as I received the wire.", "A child is reading in her bedroom."]}
|
||||
{"query": "Two woman are playing instruments; one a clarinet, the other a violin.", "pos": ["Some people are playing a tune."], "neg": ["Two women are playing a guitar and drums.", "A man is skiing down a mountain.", "The fatal dose was not taken when the murderer thought it would be.", "Person on bike", "The girl is standing, leaning against the archway.", "A group of women watch soap operas.", "No matter how old people get they never forget. "]}
|
||||
{"query": "A girl with a blue tank top sitting watching three dogs.", "pos": ["A girl is wearing blue."], "neg": ["A girl is with three cats.", "The people are watching a funeral procession.", "The child is wearing black.", "Financing is an issue for us in public schools.", "Kids at a pool.", "It is calming to be assaulted.", "I face a serious problem at eighteen years old. "]}
|
||||
{"query": "A yellow dog running along a forest path.", "pos": ["a dog is running"], "neg": ["a cat is running", "Steele did not keep her original story.", "The rule discourages people to pay their child support.", "A man in a vest sits in a car.", "Person in black clothing, with white bandanna and sunglasses waits at a bus stop.", "Neither the Globe or Mail had comments on the current state of Canada's road system. ", "The Spring Creek facility is old and outdated."]}
|
||||
{"query": "It sets out essential activities in each phase along with critical factors related to those activities.", "pos": ["Critical factors for essential activities are set out."], "neg": ["It lays out critical activities but makes no provision for critical factors related to those activities.", "People are assembled in protest.", "The state would prefer for you to do that.", "A girl sits beside a boy.", "Two males are performing.", "Nobody is jumping", "Conrad was being plotted against, to be hit on the head."]}
|
||||
{"query": "A man giving a speech in a restaurant.", "pos": ["A person gives a speech."], "neg": ["The man sits at the table and eats food.", "This is definitely not an endorsement.", "They sold their home because they were retiring and not because of the loan.", "The seal of Missouri is perfect.", "Someone is raising their hand.", "An athlete is competing in the 1500 meter swimming competition.", "Two men watching a magic show."]}
|
||||
{"query": "Indians having a gathering with coats and food and drinks.", "pos": ["A group of Indians are having a gathering with food and drinks"], "neg": ["A group of Indians are having a funeral", "It is only staged on Winter afternoons in Palma's large bullring.", "Right information can empower the legal service practices and the justice system. ", "Meanwhile, the mainland was empty of population.", "Two children is sleeping.", "a fisherman is trying to catch a monkey", "the people are in a train"]}
|
||||
{"query": "A woman with violet hair rides her bicycle outside.", "pos": ["A woman is riding her bike."], "neg": ["A woman is jogging in the park.", "The street was lined with white-painted houses.", "A group watches a movie inside.", "man at picnics cut steak", "Several chefs are sitting down and talking about food.", "The Commission notes that no significant alternatives were considered.", "We ran out of firewood and had to use pine needles for the fire."]}
|
||||
{"query": "A man pulls two women down a city street in a rickshaw.", "pos": ["A man is in a city."], "neg": ["A man is a pilot of an airplane.", "It is boring and mundane.", "The morning sunlight was shining brightly and it was warm. ", "Two people jumped off the dock.", "People watching a spaceship launch.", "Mother Teresa is an easy choice.", "It's worth being able to go at a pace you prefer."]}
|
||||
@@ -0,0 +1,43 @@
|
||||
# Pre-train
|
||||
In this example, we show how to do pre-training using retromae,
|
||||
which can improve the retrieval performance.
|
||||
|
||||
## 1. Installation
|
||||
* **with pip**
|
||||
```
|
||||
pip install -U FlagEmbedding
|
||||
```
|
||||
|
||||
* **from source**
|
||||
```
|
||||
git clone https://github.com/FlagOpen/FlagEmbedding.git
|
||||
cd FlagEmbedding/research/old-examples/pretrain
|
||||
```
|
||||
|
||||
## 2. Data format
|
||||
Train data should be a json file, where each line is a dict like this:
|
||||
```
|
||||
{"text": str}
|
||||
```
|
||||
See [toy_pretrain_data.jsonl](https://github.com/FlagOpen/FlagEmbedding/blob/master/examples/pretrain/toy_pretrain_data.jsonl) for a toy data file.
|
||||
|
||||
## 3. Train
|
||||
|
||||
```bash
|
||||
torchrun --nproc_per_node {number of gpus} \
|
||||
-m retromae_pretrain.run \
|
||||
--output_dir {path to save model} \
|
||||
--model_name_or_path BAAI/bge-large-en \
|
||||
--train_data toy_pretrain_data.jsonl \
|
||||
--learning_rate 2e-5 \
|
||||
--num_train_epochs 2 \
|
||||
--per_device_train_batch_size {batch size; set 1 for toy data} \
|
||||
--dataloader_drop_last True \
|
||||
--max_seq_length 512 \
|
||||
--logging_steps 10 \
|
||||
--dataloader_num_workers 12
|
||||
```
|
||||
|
||||
More training arguments please refer to [transformers.TrainingArguments](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments).
|
||||
After training, the encoder model will saved to `{output_dir}/encoder_model`
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataTrainingArguments:
|
||||
train_data: Optional[str] = field(
|
||||
default=None, metadata={"help": "Path to pretrain data"}
|
||||
)
|
||||
tokenizer_name: Optional[str] = field(
|
||||
default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
|
||||
)
|
||||
max_seq_length: Optional[int] = field(
|
||||
default=512,
|
||||
metadata={
|
||||
"help": "The maximum total input sequence length after tokenization. Sequences longer "
|
||||
"than this will be truncated. Default to the max input length of the model."
|
||||
},
|
||||
)
|
||||
encoder_mlm_probability: float = field(default=0.3, metadata={"help": "mask ratio for encoder"})
|
||||
decoder_mlm_probability: float = field(default=0.5, metadata={"help": "mask ratio for decoder"})
|
||||
|
||||
def __post_init__(self):
|
||||
if not os.path.exists(self.train_data):
|
||||
raise FileNotFoundError(f"cannot find file: {self.train_data}, please set a true path")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArguments:
|
||||
"""
|
||||
Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
|
||||
"""
|
||||
model_name_or_path: Optional[str] = field(
|
||||
default='bert-base-uncased',
|
||||
metadata={
|
||||
"help": "The model checkpoint for weights initialization."
|
||||
"Don't set if you want to train a model from scratch."
|
||||
},
|
||||
)
|
||||
config_name: Optional[str] = field(
|
||||
default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
|
||||
)
|
||||
@@ -0,0 +1,100 @@
|
||||
import os
|
||||
import random
|
||||
from copy import deepcopy
|
||||
from dataclasses import dataclass
|
||||
|
||||
import torch.utils.data.dataset
|
||||
from datasets import Dataset, load_dataset, concatenate_datasets
|
||||
from transformers import DataCollatorForWholeWordMask
|
||||
|
||||
from .utils import tensorize_batch
|
||||
|
||||
|
||||
class DatasetForPretraining(torch.utils.data.Dataset):
|
||||
def __init__(self, data_dir):
|
||||
if os.path.isdir(data_dir):
|
||||
datasets = []
|
||||
for file in os.listdir(data_dir):
|
||||
print(f"Loading {file}")
|
||||
file = os.path.join(data_dir, file)
|
||||
datasets.append(self.load_dataset(file))
|
||||
self.dataset = concatenate_datasets(datasets)
|
||||
else:
|
||||
print(f"Loading {data_dir}")
|
||||
self.dataset = self.load_dataset(data_dir)
|
||||
|
||||
def load_dataset(self, file):
|
||||
if file.endswith('.jsonl') or file.endswith('.json'):
|
||||
return load_dataset('json', data_files=file)['train']
|
||||
elif os.path.isdir(file):
|
||||
return Dataset.load_from_disk(file)
|
||||
else:
|
||||
raise NotImplementedError(f"Not support this file format:{file}")
|
||||
|
||||
def __getitem__(self, item):
|
||||
return self.dataset[item]['text']
|
||||
|
||||
def __len__(self):
|
||||
return len(self.dataset)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetroMAECollator(DataCollatorForWholeWordMask):
|
||||
max_seq_length: int = 512
|
||||
encoder_mlm_probability: float = 0.15
|
||||
decoder_mlm_probability: float = 0.15
|
||||
|
||||
def __call__(self, examples):
|
||||
input_ids_batch = []
|
||||
attention_mask_batch = []
|
||||
encoder_mlm_mask_batch = []
|
||||
decoder_labels_batch = []
|
||||
decoder_matrix_attention_mask_batch = []
|
||||
|
||||
for e in examples:
|
||||
|
||||
e_trunc = self.tokenizer.encode(e, max_length=self.max_seq_length, truncation=True)
|
||||
tokens = [self.tokenizer._convert_id_to_token(tid) for tid in e_trunc]
|
||||
|
||||
self.mlm_probability = self.encoder_mlm_probability
|
||||
text_encoder_mlm_mask = self._whole_word_mask(tokens)
|
||||
|
||||
self.mlm_probability = self.decoder_mlm_probability
|
||||
mask_set = []
|
||||
for _ in range(min(len(tokens), 128)):
|
||||
mask_set.append(self._whole_word_mask(tokens))
|
||||
|
||||
text_matrix_attention_mask = []
|
||||
for i in range(len(tokens)):
|
||||
idx = random.randint(0, min(len(tokens), 128) - 1)
|
||||
text_decoder_mlm_mask = deepcopy(mask_set[idx])
|
||||
text_decoder_mlm_mask[i] = 1
|
||||
text_matrix_attention_mask.append(text_decoder_mlm_mask)
|
||||
|
||||
input_ids_batch.append(torch.tensor(e_trunc))
|
||||
attention_mask_batch.append(torch.tensor([1] * len(e_trunc)))
|
||||
e_trunc[0] = -100
|
||||
e_trunc[-1] = -100
|
||||
decoder_labels_batch.append(torch.tensor(e_trunc))
|
||||
|
||||
encoder_mlm_mask_batch.append(torch.tensor(text_encoder_mlm_mask))
|
||||
decoder_matrix_attention_mask_batch.append(1 - torch.tensor(text_matrix_attention_mask))
|
||||
|
||||
input_ids_batch = tensorize_batch(input_ids_batch, self.tokenizer.pad_token_id)
|
||||
attention_mask_batch = tensorize_batch(attention_mask_batch, 0)
|
||||
origin_input_ids_batch = input_ids_batch.clone()
|
||||
encoder_mlm_mask_batch = tensorize_batch(encoder_mlm_mask_batch, 0)
|
||||
encoder_input_ids_batch, encoder_labels_batch = self.torch_mask_tokens(input_ids_batch, encoder_mlm_mask_batch)
|
||||
decoder_labels_batch = tensorize_batch(decoder_labels_batch, -100)
|
||||
matrix_attention_mask_batch = tensorize_batch(decoder_matrix_attention_mask_batch, 0)
|
||||
|
||||
batch = {
|
||||
"encoder_input_ids": encoder_input_ids_batch,
|
||||
"encoder_attention_mask": attention_mask_batch,
|
||||
"encoder_labels": encoder_labels_batch,
|
||||
"decoder_input_ids": origin_input_ids_batch,
|
||||
"decoder_attention_mask": matrix_attention_mask_batch, # [B,L,L]
|
||||
"decoder_labels": decoder_labels_batch,
|
||||
}
|
||||
|
||||
return batch
|
||||
@@ -0,0 +1,288 @@
|
||||
'''
|
||||
The codes are modified based on huggingface transformers library.
|
||||
'''
|
||||
|
||||
import math
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
import torch.utils.checkpoint
|
||||
from torch import nn
|
||||
from transformers.modeling_utils import (
|
||||
apply_chunking_to_forward,
|
||||
find_pruneable_heads_and_indices,
|
||||
prune_linear_layer,
|
||||
)
|
||||
from transformers.models.bert.modeling_bert import BertIntermediate, BertOutput, BertSelfOutput
|
||||
from transformers.utils import (
|
||||
logging,
|
||||
)
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
class BertSelfAttention(nn.Module):
|
||||
def __init__(self, config, position_embedding_type=None):
|
||||
super().__init__()
|
||||
if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
|
||||
raise ValueError(
|
||||
f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
|
||||
f"heads ({config.num_attention_heads})"
|
||||
)
|
||||
|
||||
self.num_attention_heads = config.num_attention_heads
|
||||
self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
|
||||
self.all_head_size = self.num_attention_heads * self.attention_head_size
|
||||
|
||||
self.query = nn.Linear(config.hidden_size, self.all_head_size)
|
||||
self.key = nn.Linear(config.hidden_size, self.all_head_size)
|
||||
self.value = nn.Linear(config.hidden_size, self.all_head_size)
|
||||
|
||||
self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
|
||||
self.position_embedding_type = position_embedding_type or getattr(
|
||||
config, "position_embedding_type", "absolute"
|
||||
)
|
||||
if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
|
||||
self.max_position_embeddings = config.max_position_embeddings
|
||||
self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)
|
||||
|
||||
self.is_decoder = config.is_decoder
|
||||
|
||||
def transpose_for_scores(self, x):
|
||||
new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
|
||||
x = x.view(new_x_shape)
|
||||
return x.permute(0, 2, 1, 3)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
attention_mask: Optional[torch.FloatTensor] = None,
|
||||
head_mask: Optional[torch.FloatTensor] = None,
|
||||
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
||||
encoder_attention_mask: Optional[torch.FloatTensor] = None,
|
||||
past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
|
||||
output_attentions: Optional[bool] = False,
|
||||
) -> Tuple[torch.Tensor]:
|
||||
mixed_query_layer = self.query(query)
|
||||
|
||||
# If this is instantiated as a cross-attention module, the keys
|
||||
# and values come from an encoder; the attention mask needs to be
|
||||
# such that the encoder's padding tokens are not attended to.
|
||||
is_cross_attention = encoder_hidden_states is not None
|
||||
|
||||
if is_cross_attention and past_key_value is not None:
|
||||
# reuse k,v, cross_attentions
|
||||
key_layer = past_key_value[0]
|
||||
value_layer = past_key_value[1]
|
||||
attention_mask = encoder_attention_mask
|
||||
elif is_cross_attention:
|
||||
key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
|
||||
value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
|
||||
attention_mask = encoder_attention_mask
|
||||
elif past_key_value is not None:
|
||||
key_layer = self.transpose_for_scores(self.key(key))
|
||||
value_layer = self.transpose_for_scores(self.value(value))
|
||||
key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
|
||||
value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
|
||||
else:
|
||||
key_layer = self.transpose_for_scores(self.key(key))
|
||||
value_layer = self.transpose_for_scores(self.value(value))
|
||||
|
||||
query_layer = self.transpose_for_scores(mixed_query_layer)
|
||||
|
||||
if self.is_decoder:
|
||||
# if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
|
||||
# Further calls to cross_attention layer can then reuse all cross-attention
|
||||
# key/value_states (first "if" case)
|
||||
# if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
|
||||
# all previous decoder key/value_states. Further calls to uni-directional self-attention
|
||||
# can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
|
||||
# if encoder bi-directional self-attention `past_key_value` is always `None`
|
||||
past_key_value = (key_layer, value_layer)
|
||||
|
||||
# Take the dot product between "query" and "key" to get the raw attention scores.
|
||||
attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
|
||||
|
||||
if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
|
||||
seq_length = query.size()[1]
|
||||
position_ids_l = torch.arange(seq_length, dtype=torch.long, device=query.device).view(-1, 1)
|
||||
position_ids_r = torch.arange(seq_length, dtype=torch.long, device=query.device).view(1, -1)
|
||||
distance = position_ids_l - position_ids_r
|
||||
positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
|
||||
positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility
|
||||
|
||||
if self.position_embedding_type == "relative_key":
|
||||
relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
|
||||
attention_scores = attention_scores + relative_position_scores
|
||||
elif self.position_embedding_type == "relative_key_query":
|
||||
relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
|
||||
relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
|
||||
attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key
|
||||
|
||||
attention_scores = attention_scores / math.sqrt(self.attention_head_size)
|
||||
if attention_mask is not None:
|
||||
# Apply the attention mask is (precomputed for all layers in BertModel forward() function)
|
||||
attention_scores = attention_scores + attention_mask
|
||||
|
||||
# Normalize the attention scores to probabilities.
|
||||
attention_probs = nn.functional.softmax(attention_scores, dim=-1)
|
||||
|
||||
# This is actually dropping out entire tokens to attend to, which might
|
||||
# seem a bit unusual, but is taken from the original Transformer paper.
|
||||
attention_probs = self.dropout(attention_probs)
|
||||
|
||||
# Mask heads if we want to
|
||||
if head_mask is not None:
|
||||
attention_probs = attention_probs * head_mask
|
||||
|
||||
context_layer = torch.matmul(attention_probs, value_layer)
|
||||
|
||||
context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
|
||||
new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
|
||||
context_layer = context_layer.view(new_context_layer_shape)
|
||||
|
||||
outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
|
||||
|
||||
if self.is_decoder:
|
||||
outputs = outputs + (past_key_value,)
|
||||
return outputs
|
||||
|
||||
|
||||
class BertAttention(nn.Module):
|
||||
def __init__(self, config, position_embedding_type=None):
|
||||
super().__init__()
|
||||
self.self = BertSelfAttention(config, position_embedding_type=position_embedding_type)
|
||||
self.output = BertSelfOutput(config)
|
||||
self.pruned_heads = set()
|
||||
|
||||
def prune_heads(self, heads):
|
||||
if len(heads) == 0:
|
||||
return
|
||||
heads, index = find_pruneable_heads_and_indices(
|
||||
heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
|
||||
)
|
||||
|
||||
# Prune linear layers
|
||||
self.self.query = prune_linear_layer(self.self.query, index)
|
||||
self.self.key = prune_linear_layer(self.self.key, index)
|
||||
self.self.value = prune_linear_layer(self.self.value, index)
|
||||
self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
|
||||
|
||||
# Update hyper params and store pruned heads
|
||||
self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
|
||||
self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
|
||||
self.pruned_heads = self.pruned_heads.union(heads)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attention_mask: Optional[torch.FloatTensor] = None,
|
||||
head_mask: Optional[torch.FloatTensor] = None,
|
||||
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
||||
encoder_attention_mask: Optional[torch.FloatTensor] = None,
|
||||
past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
|
||||
output_attentions: Optional[bool] = False,
|
||||
) -> Tuple[torch.Tensor]:
|
||||
self_outputs = self.self(
|
||||
query, key, value,
|
||||
attention_mask,
|
||||
head_mask,
|
||||
encoder_hidden_states,
|
||||
encoder_attention_mask,
|
||||
past_key_value,
|
||||
output_attentions,
|
||||
)
|
||||
attention_output = self.output(self_outputs[0], query)
|
||||
outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
|
||||
return outputs
|
||||
|
||||
|
||||
class BertLayerForDecoder(nn.Module):
|
||||
def __init__(self, config):
|
||||
super().__init__()
|
||||
self.chunk_size_feed_forward = config.chunk_size_feed_forward
|
||||
self.seq_len_dim = 1
|
||||
self.attention = BertAttention(config)
|
||||
self.is_decoder = config.is_decoder
|
||||
self.add_cross_attention = config.add_cross_attention
|
||||
if self.add_cross_attention:
|
||||
if not self.is_decoder:
|
||||
raise ValueError(f"{self} should be used as a decoder model if cross attention is added")
|
||||
self.crossattention = BertAttention(config, position_embedding_type="absolute")
|
||||
self.intermediate = BertIntermediate(config)
|
||||
self.output = BertOutput(config)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
query: torch.Tensor,
|
||||
key: torch.Tensor,
|
||||
value: torch.Tensor,
|
||||
attention_mask: Optional[torch.FloatTensor] = None,
|
||||
head_mask: Optional[torch.FloatTensor] = None,
|
||||
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
||||
encoder_attention_mask: Optional[torch.FloatTensor] = None,
|
||||
past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
|
||||
output_attentions: Optional[bool] = False,
|
||||
) -> Tuple[torch.Tensor]:
|
||||
# decoder uni-directional self-attention cached key/values tuple is at positions 1,2
|
||||
self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
|
||||
self_attention_outputs = self.attention(
|
||||
query, key, value,
|
||||
attention_mask,
|
||||
head_mask,
|
||||
output_attentions=output_attentions,
|
||||
past_key_value=self_attn_past_key_value,
|
||||
)
|
||||
attention_output = self_attention_outputs[0]
|
||||
|
||||
# if decoder, the last output is tuple of self-attn cache
|
||||
if self.is_decoder:
|
||||
outputs = self_attention_outputs[1:-1]
|
||||
present_key_value = self_attention_outputs[-1]
|
||||
else:
|
||||
outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
|
||||
|
||||
cross_attn_present_key_value = None
|
||||
if self.is_decoder and encoder_hidden_states is not None:
|
||||
if not hasattr(self, "crossattention"):
|
||||
raise ValueError(
|
||||
f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`"
|
||||
)
|
||||
|
||||
# cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
|
||||
cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
|
||||
cross_attention_outputs = self.crossattention(
|
||||
attention_output,
|
||||
attention_mask,
|
||||
head_mask,
|
||||
encoder_hidden_states,
|
||||
encoder_attention_mask,
|
||||
cross_attn_past_key_value,
|
||||
output_attentions,
|
||||
)
|
||||
attention_output = cross_attention_outputs[0]
|
||||
outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights
|
||||
|
||||
# add cross-attn cache to positions 3,4 of present_key_value tuple
|
||||
cross_attn_present_key_value = cross_attention_outputs[-1]
|
||||
present_key_value = present_key_value + cross_attn_present_key_value
|
||||
|
||||
layer_output = apply_chunking_to_forward(
|
||||
self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
|
||||
)
|
||||
outputs = (layer_output,) + outputs
|
||||
|
||||
# if decoder, return the attn key/values as the last output
|
||||
if self.is_decoder:
|
||||
outputs = outputs + (present_key_value,)
|
||||
|
||||
return outputs
|
||||
|
||||
def feed_forward_chunk(self, attention_output):
|
||||
intermediate_output = self.intermediate(attention_output)
|
||||
layer_output = self.output(intermediate_output, attention_output)
|
||||
return layer_output
|
||||
@@ -0,0 +1,102 @@
|
||||
import logging
|
||||
import os
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from transformers import BertForMaskedLM, AutoModelForMaskedLM
|
||||
from transformers.modeling_outputs import MaskedLMOutput
|
||||
|
||||
from .arguments import ModelArguments
|
||||
from .enhancedDecoder import BertLayerForDecoder
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RetroMAEForPretraining(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
bert: BertForMaskedLM,
|
||||
model_args: ModelArguments,
|
||||
):
|
||||
super(RetroMAEForPretraining, self).__init__()
|
||||
self.lm = bert
|
||||
|
||||
if hasattr(self.lm, 'bert'):
|
||||
self.decoder_embeddings = self.lm.bert.embeddings
|
||||
elif hasattr(self.lm, 'roberta'):
|
||||
self.decoder_embeddings = self.lm.roberta.embeddings
|
||||
else:
|
||||
self.decoder_embeddings = self.lm.bert.embeddings
|
||||
|
||||
self.c_head = BertLayerForDecoder(bert.config)
|
||||
self.c_head.apply(self.lm._init_weights)
|
||||
|
||||
self.cross_entropy = nn.CrossEntropyLoss()
|
||||
|
||||
self.model_args = model_args
|
||||
|
||||
def gradient_checkpointing_enable(self, **kwargs):
|
||||
self.lm.gradient_checkpointing_enable(**kwargs)
|
||||
|
||||
def forward(self,
|
||||
encoder_input_ids, encoder_attention_mask, encoder_labels,
|
||||
decoder_input_ids, decoder_attention_mask, decoder_labels):
|
||||
|
||||
lm_out: MaskedLMOutput = self.lm(
|
||||
encoder_input_ids, encoder_attention_mask,
|
||||
labels=encoder_labels,
|
||||
output_hidden_states=True,
|
||||
return_dict=True
|
||||
)
|
||||
cls_hiddens = lm_out.hidden_states[-1][:, :1] # B 1 D
|
||||
|
||||
decoder_embedding_output = self.decoder_embeddings(input_ids=decoder_input_ids)
|
||||
hiddens = torch.cat([cls_hiddens, decoder_embedding_output[:, 1:]], dim=1)
|
||||
|
||||
# decoder_position_ids = self.lm.bert.embeddings.position_ids[:, :decoder_input_ids.size(1)]
|
||||
# decoder_position_embeddings = self.lm.bert.embeddings.position_embeddings(decoder_position_ids) # B L D
|
||||
# query = decoder_position_embeddings + cls_hiddens
|
||||
|
||||
cls_hiddens = cls_hiddens.expand(hiddens.size(0), hiddens.size(1), hiddens.size(2))
|
||||
query = self.decoder_embeddings(inputs_embeds=cls_hiddens)
|
||||
|
||||
matrix_attention_mask = self.lm.get_extended_attention_mask(
|
||||
decoder_attention_mask,
|
||||
decoder_attention_mask.shape,
|
||||
decoder_attention_mask.device
|
||||
)
|
||||
|
||||
hiddens = self.c_head(query=query,
|
||||
key=hiddens,
|
||||
value=hiddens,
|
||||
attention_mask=matrix_attention_mask)[0]
|
||||
pred_scores, loss = self.mlm_loss(hiddens, decoder_labels)
|
||||
|
||||
return (loss + lm_out.loss,)
|
||||
|
||||
def mlm_loss(self, hiddens, labels):
|
||||
if hasattr(self.lm, 'cls'):
|
||||
pred_scores = self.lm.cls(hiddens)
|
||||
elif hasattr(self.lm, 'lm_head'):
|
||||
pred_scores = self.lm.lm_head(hiddens)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
masked_lm_loss = self.cross_entropy(
|
||||
pred_scores.view(-1, self.lm.config.vocab_size),
|
||||
labels.view(-1)
|
||||
)
|
||||
return pred_scores, masked_lm_loss
|
||||
|
||||
def save_pretrained(self, output_dir: str):
|
||||
self.lm.save_pretrained(os.path.join(output_dir, "encoder_model"))
|
||||
torch.save(self.state_dict(), os.path.join(output_dir, 'pytorch_model.bin'))
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls, model_args: ModelArguments,
|
||||
*args, **kwargs
|
||||
):
|
||||
hf_model = AutoModelForMaskedLM.from_pretrained(*args, **kwargs)
|
||||
model = cls(hf_model, model_args)
|
||||
return model
|
||||
@@ -0,0 +1,128 @@
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
import transformers
|
||||
from transformers import (
|
||||
AutoTokenizer,
|
||||
BertForMaskedLM,
|
||||
AutoConfig,
|
||||
HfArgumentParser, set_seed, )
|
||||
from transformers import (
|
||||
TrainerCallback,
|
||||
TrainingArguments,
|
||||
TrainerState,
|
||||
TrainerControl
|
||||
)
|
||||
from transformers.trainer_utils import is_main_process
|
||||
|
||||
from .arguments import DataTrainingArguments, ModelArguments
|
||||
from .data import DatasetForPretraining, RetroMAECollator
|
||||
from .modeling import RetroMAEForPretraining
|
||||
from .trainer import PreTrainer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TrainerCallbackForSaving(TrainerCallback):
|
||||
def on_epoch_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
|
||||
"""
|
||||
Event called at the end of an epoch.
|
||||
"""
|
||||
control.should_save = True
|
||||
|
||||
|
||||
def main():
|
||||
# See all possible arguments in src/transformers/training_args.py
|
||||
# or by passing the --help flag to this script.
|
||||
# We now keep distinct sets of args, for a cleaner separation of concerns.
|
||||
|
||||
parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
|
||||
if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
|
||||
# If we pass only one argument to the script and it's the path to a json file,
|
||||
# let's parse it to get our arguments.
|
||||
model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
|
||||
else:
|
||||
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
|
||||
|
||||
if (
|
||||
os.path.exists(training_args.output_dir)
|
||||
and os.listdir(training_args.output_dir)
|
||||
and training_args.do_train
|
||||
and not training_args.overwrite_output_dir
|
||||
):
|
||||
raise ValueError(
|
||||
f"Output directory ({training_args.output_dir}) already exists and is not empty."
|
||||
"Use --overwrite_output_dir to overcome."
|
||||
)
|
||||
|
||||
model_args: ModelArguments
|
||||
data_args: DataTrainingArguments
|
||||
training_args: TrainingArguments
|
||||
|
||||
training_args.remove_unused_columns = False
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
datefmt="%m/%d/%Y %H:%M:%S",
|
||||
level=logging.INFO if is_main_process(training_args.local_rank) else logging.WARN,
|
||||
)
|
||||
|
||||
# Log on each process the small summary:
|
||||
logger.warning(
|
||||
f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
|
||||
+ f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
|
||||
)
|
||||
# Set the verbosity to info of the Transformers logger (on main process only):
|
||||
if is_main_process(training_args.local_rank):
|
||||
transformers.utils.logging.set_verbosity_info()
|
||||
transformers.utils.logging.enable_default_handler()
|
||||
transformers.utils.logging.enable_explicit_format()
|
||||
if training_args.local_rank in (0, -1):
|
||||
logger.info("Training/evaluation parameters %s", training_args)
|
||||
logger.info("Model parameters %s", model_args)
|
||||
logger.info("Data parameters %s", data_args)
|
||||
|
||||
set_seed(training_args.seed)
|
||||
|
||||
model_class = RetroMAEForPretraining
|
||||
collator_class = RetroMAECollator
|
||||
|
||||
if model_args.model_name_or_path:
|
||||
model = model_class.from_pretrained(model_args, model_args.model_name_or_path)
|
||||
logger.info(f"------Load model from {model_args.model_name_or_path}------")
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path)
|
||||
elif model_args.config_name:
|
||||
config = AutoConfig.from_pretrained(model_args.config_name)
|
||||
bert = BertForMaskedLM(config)
|
||||
model = model_class(bert, model_args)
|
||||
logger.info("------Init the model------")
|
||||
tokenizer = AutoTokenizer.from_pretrained(data_args.tokenizer_name)
|
||||
else:
|
||||
raise ValueError("You must provide the model_name_or_path or config_name")
|
||||
|
||||
dataset = DatasetForPretraining(data_args.train_data)
|
||||
|
||||
data_collator = collator_class(tokenizer,
|
||||
encoder_mlm_probability=data_args.encoder_mlm_probability,
|
||||
decoder_mlm_probability=data_args.decoder_mlm_probability,
|
||||
max_seq_length=data_args.max_seq_length)
|
||||
|
||||
# Initialize our Trainer
|
||||
trainer = PreTrainer(
|
||||
model=model,
|
||||
args=training_args,
|
||||
train_dataset=dataset,
|
||||
data_collator=data_collator,
|
||||
tokenizer=tokenizer
|
||||
)
|
||||
trainer.add_callback(TrainerCallbackForSaving())
|
||||
|
||||
# # Training
|
||||
trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
|
||||
trainer.save_model() # Saves the tokenizer too for easy upload
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,46 @@
|
||||
import logging
|
||||
import os
|
||||
from typing import Dict, Optional
|
||||
|
||||
import torch
|
||||
from transformers import Trainer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PreTrainer(Trainer):
|
||||
def log(self, logs: Dict[str, float]) -> None:
|
||||
"""
|
||||
Log `logs` on the various objects watching training.
|
||||
|
||||
Subclass and override this method to inject custom behavior.
|
||||
|
||||
Args:
|
||||
logs (`Dict[str, float]`):
|
||||
The values to log.
|
||||
"""
|
||||
logs["step"] = self.state.global_step
|
||||
if self.state.epoch is not None:
|
||||
logs["epoch"] = round(self.state.epoch, 2)
|
||||
|
||||
output = {**logs, **{"step": self.state.global_step}}
|
||||
self.state.log_history.append(output)
|
||||
self.control = self.callback_handler.on_log(self.args, self.state, self.control, logs)
|
||||
|
||||
def _save(self, output_dir: Optional[str] = None, state_dict=None):
|
||||
output_dir = output_dir if output_dir is not None else self.args.output_dir
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
logger.info(f"Saving model checkpoint to {output_dir}")
|
||||
# Save a trained model and configuration using `save_pretrained()`.
|
||||
# They can then be reloaded using `from_pretrained()`
|
||||
if not hasattr(self.model, 'save_pretrained'):
|
||||
logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.")
|
||||
state_dict = self.model.state_dict()
|
||||
torch.save(state_dict, os.path.join(output_dir, "pytorch_model.bin"))
|
||||
else:
|
||||
self.model.save_pretrained(output_dir)
|
||||
if self.tokenizer is not None:
|
||||
self.tokenizer.save_pretrained(os.path.join(output_dir, "encoder_model"))
|
||||
|
||||
# Good practice: save your training arguments together with the trained model
|
||||
torch.save(self.args, os.path.join(output_dir, "training_args.bin"))
|
||||
@@ -0,0 +1,32 @@
|
||||
from typing import List
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def tensorize_batch(sequences: List[torch.Tensor], padding_value, align_right=False) -> torch.Tensor:
|
||||
if len(sequences[0].size()) == 1:
|
||||
max_len_1 = max([s.size(0) for s in sequences])
|
||||
out_dims = (len(sequences), max_len_1)
|
||||
out_tensor = sequences[0].new_full(out_dims, padding_value)
|
||||
for i, tensor in enumerate(sequences):
|
||||
length_1 = tensor.size(0)
|
||||
if align_right:
|
||||
out_tensor[i, -length_1:] = tensor
|
||||
else:
|
||||
out_tensor[i, :length_1] = tensor
|
||||
return out_tensor
|
||||
elif len(sequences[0].size()) == 2:
|
||||
max_len_1 = max([s.size(0) for s in sequences])
|
||||
max_len_2 = max([s.size(1) for s in sequences])
|
||||
out_dims = (len(sequences), max_len_1, max_len_2)
|
||||
out_tensor = sequences[0].new_full(out_dims, padding_value)
|
||||
for i, tensor in enumerate(sequences):
|
||||
length_1 = tensor.size(0)
|
||||
length_2 = tensor.size(1)
|
||||
if align_right:
|
||||
out_tensor[i, -length_1:, :length_2] = tensor
|
||||
else:
|
||||
out_tensor[i, :length_1, :length_2] = tensor
|
||||
return out_tensor
|
||||
else:
|
||||
raise
|
||||
@@ -0,0 +1,10 @@
|
||||
{"text": "RetroMAE can provide a strong initialization of dense retriever; after fine-tuned with in-domain data, it gives rise to a high-quality supervised retrieval performance in the corresponding scenario. Besides, It substantially improves the pre-trained model's transferability, which helps to result in superior zero-shot performances on out-of-domain datasets."}
|
||||
{"text": "RetroMAE can provide a strong initialization of dense retriever; after fine-tuned with in-domain data, it gives rise to a high-quality supervised retrieval performance in the corresponding scenario. Besides, It substantially improves the pre-trained model's transferability, which helps to result in superior zero-shot performances on out-of-domain datasets."}
|
||||
{"text": "RetroMAE can provide a strong initialization of dense retriever; after fine-tuned with in-domain data, it gives rise to a high-quality supervised retrieval performance in the corresponding scenario. Besides, It substantially improves the pre-trained model's transferability, which helps to result in superior zero-shot performances on out-of-domain datasets."}
|
||||
{"text": "RetroMAE can provide a strong initialization of dense retriever; after fine-tuned with in-domain data, it gives rise to a high-quality supervised retrieval performance in the corresponding scenario. Besides, It substantially improves the pre-trained model's transferability, which helps to result in superior zero-shot performances on out-of-domain datasets."}
|
||||
{"text": "RetroMAE can provide a strong initialization of dense retriever; after fine-tuned with in-domain data, it gives rise to a high-quality supervised retrieval performance in the corresponding scenario. Besides, It substantially improves the pre-trained model's transferability, which helps to result in superior zero-shot performances on out-of-domain datasets."}
|
||||
{"text": "RetroMAE can provide a strong initialization of dense retriever; after fine-tuned with in-domain data, it gives rise to a high-quality supervised retrieval performance in the corresponding scenario. Besides, It substantially improves the pre-trained model's transferability, which helps to result in superior zero-shot performances on out-of-domain datasets."}
|
||||
{"text": "RetroMAE can provide a strong initialization of dense retriever; after fine-tuned with in-domain data, it gives rise to a high-quality supervised retrieval performance in the corresponding scenario. Besides, It substantially improves the pre-trained model's transferability, which helps to result in superior zero-shot performances on out-of-domain datasets."}
|
||||
{"text": "RetroMAE can provide a strong initialization of dense retriever; after fine-tuned with in-domain data, it gives rise to a high-quality supervised retrieval performance in the corresponding scenario. Besides, It substantially improves the pre-trained model's transferability, which helps to result in superior zero-shot performances on out-of-domain datasets."}
|
||||
{"text": "RetroMAE can provide a strong initialization of dense retriever; after fine-tuned with in-domain data, it gives rise to a high-quality supervised retrieval performance in the corresponding scenario. Besides, It substantially improves the pre-trained model's transferability, which helps to result in superior zero-shot performances on out-of-domain datasets."}
|
||||
{"text": "RetroMAE can provide a strong initialization of dense retriever; after fine-tuned with in-domain data, it gives rise to a high-quality supervised retrieval performance in the corresponding scenario. Besides, It substantially improves the pre-trained model's transferability, which helps to result in superior zero-shot performances on out-of-domain datasets."}
|
||||
@@ -0,0 +1,108 @@
|
||||
# Finetune cross-encoder
|
||||
In this example, we show how to finetune the cross-encoder reranker with your data.
|
||||
|
||||
## 1. Installation
|
||||
```
|
||||
git clone https://github.com/FlagOpen/FlagEmbedding.git
|
||||
cd research/reranker
|
||||
pip install .
|
||||
```
|
||||
|
||||
## 2. Data format
|
||||
|
||||
The data format for reranker is the same as [embedding fine-tune](https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/finetune/embedder#2-data-format).
|
||||
Besides, we strongly suggest to [mine hard negatives](https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/finetune/reranker#hard-negatives) to fine-tune reranker.
|
||||
|
||||
|
||||
## 3. Train
|
||||
|
||||
```
|
||||
torchrun --nproc_per_node {number of gpus} \
|
||||
-m run \
|
||||
--output_dir {path to save model} \
|
||||
--model_name_or_path BAAI/bge-reranker-base \
|
||||
--train_data ./toy_finetune_data.jsonl \
|
||||
--learning_rate 6e-5 \
|
||||
--fp16 \
|
||||
--num_train_epochs 5 \
|
||||
--per_device_train_batch_size {batch size; set 1 for toy data} \
|
||||
--gradient_accumulation_steps 4 \
|
||||
--dataloader_drop_last True \
|
||||
--train_group_size 16 \
|
||||
--max_len 512 \
|
||||
--weight_decay 0.01 \
|
||||
--logging_steps 10
|
||||
```
|
||||
|
||||
**some important arguments**:
|
||||
- `per_device_train_batch_size`: batch size in training.
|
||||
- `train_group_size`: the number of positive and negatives for a query in training.
|
||||
There are always one positive, so this argument will control the number of negatives (#negatives=train_group_size-1).
|
||||
Noted that the number of negatives should not be larger than the numbers of negatives in data `"neg":List[str]`.
|
||||
Besides the negatives in this group, the in-batch negatives also will be used in fine-tuning.
|
||||
|
||||
More training arguments please refer to [transformers.TrainingArguments](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments)
|
||||
|
||||
|
||||
### 4. Model merging via [LM-Cocktail](https://github.com/FlagOpen/FlagEmbedding/tree/master/research/LM_Cocktail) [optional]
|
||||
|
||||
For more details please refer to [LM-Cocktail](https://github.com/FlagOpen/FlagEmbedding/tree/master/research/LM_Cocktail).
|
||||
|
||||
Fine-tuning the base bge model can improve its performance on target task,
|
||||
but maybe lead to severe degeneration of model’s general capabilities
|
||||
beyond the targeted domain (e.g., lower performance on c-mteb tasks).
|
||||
By merging the fine-tuned model and the base model,
|
||||
LM-Cocktail can significantly enhance performance in downstream task
|
||||
while maintaining performance in other unrelated tasks.
|
||||
|
||||
```python
|
||||
from LM_Cocktail import mix_models, mix_models_with_data
|
||||
|
||||
# Mix fine-tuned model and base model; then save it to output_path: ./mixed_model_1
|
||||
model = mix_models(
|
||||
model_names_or_paths=["BAAI/bge-reranker-base", "your_fine-tuned_model"],
|
||||
model_type='reranker',
|
||||
weights=[0.5, 0.5], # you can change the weights to get a better trade-off.
|
||||
output_path='./mixed_model_1')
|
||||
```
|
||||
|
||||
|
||||
|
||||
|
||||
### 5. Load your model
|
||||
|
||||
#### Using FlagEmbedding
|
||||
|
||||
```python
|
||||
from FlagEmbedding import FlagReranker
|
||||
reranker = FlagReranker('BAAI/bge-reranker-base', use_fp16=True) #use fp16 can speed up computing
|
||||
|
||||
score = reranker.compute_score(['query', 'passage'])
|
||||
print(score)
|
||||
|
||||
scores = reranker.compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']])
|
||||
print(scores)
|
||||
```
|
||||
|
||||
|
||||
#### Using Huggingface transformers
|
||||
|
||||
```python
|
||||
import torch
|
||||
from transformers import AutoModelForSequenceClassification, AutoTokenizer, BatchEncoding, PreTrainedTokenizerFast
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-base')
|
||||
model = AutoModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-base')
|
||||
model.eval()
|
||||
|
||||
pairs = [['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']]
|
||||
with torch.no_grad():
|
||||
inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
|
||||
scores = model(**inputs, return_dict=True).logits.view(-1, ).float()
|
||||
print(scores)
|
||||
```
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
{
|
||||
"fp16": {
|
||||
"enabled": "auto",
|
||||
"loss_scale": 0,
|
||||
"loss_scale_window": 1000,
|
||||
"initial_scale_power": 12,
|
||||
"hysteresis": 2,
|
||||
"min_loss_scale": 1
|
||||
},
|
||||
|
||||
"bf16": {
|
||||
"enabled": "auto"
|
||||
},
|
||||
|
||||
"optimizer": {
|
||||
"type": "AdamW",
|
||||
"params": {
|
||||
"lr": "auto",
|
||||
"betas": "auto",
|
||||
"eps": "auto",
|
||||
"weight_decay": "auto"
|
||||
}
|
||||
},
|
||||
|
||||
"scheduler": {
|
||||
"type": "WarmupDecayLR",
|
||||
"params": {
|
||||
"warmup_min_lr": "auto",
|
||||
"warmup_max_lr": "auto",
|
||||
"warmup_num_steps": "auto",
|
||||
"total_num_steps": "auto"
|
||||
}
|
||||
},
|
||||
|
||||
"zero_optimization": {
|
||||
"stage": 0
|
||||
},
|
||||
|
||||
"gradient_accumulation_steps": "auto",
|
||||
"gradient_clipping": "auto",
|
||||
"steps_per_print": 100,
|
||||
"train_batch_size": "auto",
|
||||
"train_micro_batch_size_per_gpu": "auto",
|
||||
"wall_clock_breakdown": false
|
||||
}
|
||||
@@ -0,0 +1,10 @@
|
||||
{"query": "Five women walk along a beach wearing flip-flops.", "pos": ["Some women with flip-flops on, are walking along the beach"], "neg": ["The 4 women are sitting on the beach.", "There was a reform in 1996.", "She's not going to court to clear her record.", "The man is talking about hawaii.", "A woman is standing outside.", "The battle was over. ", "A group of people plays volleyball."]}
|
||||
{"query": "A woman standing on a high cliff on one leg looking over a river.", "pos": ["A woman is standing on a cliff."], "neg": ["A woman sits on a chair.", "George Bush told the Republicans there was no way he would let them even consider this foolish idea, against his top advisors advice.", "The family was falling apart.", "no one showed up to the meeting", "A boy is sitting outside playing in the sand.", "Ended as soon as I received the wire.", "A child is reading in her bedroom."]}
|
||||
{"query": "Two woman are playing instruments; one a clarinet, the other a violin.", "pos": ["Some people are playing a tune."], "neg": ["Two women are playing a guitar and drums.", "A man is skiing down a mountain.", "The fatal dose was not taken when the murderer thought it would be.", "Person on bike", "The girl is standing, leaning against the archway.", "A group of women watch soap operas.", "No matter how old people get they never forget. "]}
|
||||
{"query": "A girl with a blue tank top sitting watching three dogs.", "pos": ["A girl is wearing blue."], "neg": ["A girl is with three cats.", "The people are watching a funeral procession.", "The child is wearing black.", "Financing is an issue for us in public schools.", "Kids at a pool.", "It is calming to be assaulted.", "I face a serious problem at eighteen years old. "]}
|
||||
{"query": "A yellow dog running along a forest path.", "pos": ["a dog is running"], "neg": ["a cat is running", "Steele did not keep her original story.", "The rule discourages people to pay their child support.", "A man in a vest sits in a car.", "Person in black clothing, with white bandanna and sunglasses waits at a bus stop.", "Neither the Globe or Mail had comments on the current state of Canada's road system. ", "The Spring Creek facility is old and outdated."]}
|
||||
{"query": "It sets out essential activities in each phase along with critical factors related to those activities.", "pos": ["Critical factors for essential activities are set out."], "neg": ["It lays out critical activities but makes no provision for critical factors related to those activities.", "People are assembled in protest.", "The state would prefer for you to do that.", "A girl sits beside a boy.", "Two males are performing.", "Nobody is jumping", "Conrad was being plotted against, to be hit on the head."]}
|
||||
{"query": "A man giving a speech in a restaurant.", "pos": ["A person gives a speech."], "neg": ["The man sits at the table and eats food.", "This is definitely not an endorsement.", "They sold their home because they were retiring and not because of the loan.", "The seal of Missouri is perfect.", "Someone is raising their hand.", "An athlete is competing in the 1500 meter swimming competition.", "Two men watching a magic show."]}
|
||||
{"query": "Indians having a gathering with coats and food and drinks.", "pos": ["A group of Indians are having a gathering with food and drinks"], "neg": ["A group of Indians are having a funeral", "It is only staged on Winter afternoons in Palma's large bullring.", "Right information can empower the legal service practices and the justice system. ", "Meanwhile, the mainland was empty of population.", "Two children is sleeping.", "a fisherman is trying to catch a monkey", "the people are in a train"]}
|
||||
{"query": "A woman with violet hair rides her bicycle outside.", "pos": ["A woman is riding her bike."], "neg": ["A woman is jogging in the park.", "The street was lined with white-painted houses.", "A group watches a movie inside.", "man at picnics cut steak", "Several chefs are sitting down and talking about food.", "The Commission notes that no significant alternatives were considered.", "We ran out of firewood and had to use pine needles for the fire."]}
|
||||
{"query": "A man pulls two women down a city street in a rickshaw.", "pos": ["A man is in a city."], "neg": ["A man is a pilot of an airplane.", "It is boring and mundane.", "The morning sunlight was shining brightly and it was warm. ", "Two people jumped off the dock.", "People watching a spaceship launch.", "Mother Teresa is an easy choice.", "It's worth being able to go at a pace you prefer."]}
|
||||
@@ -0,0 +1,16 @@
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArguments:
|
||||
model_name_or_path: str = field(
|
||||
default='BAAI/bge-large-zh-noinstruct',
|
||||
metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataArguments:
|
||||
data_path: str = field(
|
||||
default='./data', metadata={"help": "Path to wikipedia-22-12"}
|
||||
)
|
||||
@@ -0,0 +1,112 @@
|
||||
import json
|
||||
import os
|
||||
import subprocess
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from arguments import ModelArguments, DataArguments
|
||||
from datasets import load_dataset
|
||||
from torch.utils.data import Dataset, SequentialSampler
|
||||
from torch_geometric.data import DataLoader
|
||||
from tqdm import tqdm
|
||||
from transformers import AutoTokenizer, HfArgumentParser, is_torch_npu_available
|
||||
from transformers import PreTrainedTokenizer, AutoModel
|
||||
|
||||
|
||||
class EmbDataset(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: PreTrainedTokenizer,
|
||||
path: str
|
||||
):
|
||||
self.tokenizer = tokenizer
|
||||
with open(path, 'r') as f:
|
||||
self.data = json.load(f)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data)
|
||||
|
||||
def __getitem__(self, item):
|
||||
sentences = self.data[item]['contents']
|
||||
batch_dict = self.tokenizer(sentences, max_length=512, padding=True, truncation=True, return_tensors='pt')
|
||||
attention_mask = batch_dict['attention_mask'][0].tolist() + [0] * (512 - len(batch_dict['attention_mask'][0]))
|
||||
token_type_ids = batch_dict['token_type_ids'][0].tolist() + [0] * (512 - len(batch_dict['token_type_ids'][0]))
|
||||
input_ids = batch_dict['input_ids'][0].tolist() + [0] * (512 - len(batch_dict['token_type_ids'][0]))
|
||||
|
||||
return torch.LongTensor(input_ids), torch.LongTensor(token_type_ids), torch.LongTensor(attention_mask)
|
||||
|
||||
|
||||
def inference(json_path, emb_path, model_path):
|
||||
if torch.cuda.is_available():
|
||||
device = torch.device("cuda")
|
||||
elif is_torch_npu_available():
|
||||
device = torch.device("npu")
|
||||
else:
|
||||
device = torch.device("cpu")
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
||||
|
||||
model = AutoModel.from_pretrained(model_path).to(device)
|
||||
model = torch.nn.parallel.DataParallel(model)
|
||||
|
||||
dataset = EmbDataset(tokenizer, json_path)
|
||||
loader = DataLoader(dataset=dataset, batch_size=2048, sampler=SequentialSampler(dataset), shuffle=False,
|
||||
drop_last=False, num_workers=16)
|
||||
|
||||
model.eval()
|
||||
existing_data = []
|
||||
for step, data in enumerate(tqdm(loader, total=len(loader))):
|
||||
input_ids, token_type_ids, attention_mask = data
|
||||
input_ids = input_ids.to(device)
|
||||
token_type_ids = token_type_ids.to(device)
|
||||
attention_mask = attention_mask.to(device)
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = model(input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask)
|
||||
batch_vecs = outputs[0][:, 0]
|
||||
batch_vecs = torch.nn.functional.normalize(batch_vecs, p=2, dim=-1).detach().cpu().numpy()
|
||||
|
||||
existing_data.append(batch_vecs)
|
||||
|
||||
np.save(emb_path, np.concatenate(existing_data))
|
||||
|
||||
|
||||
def build_bm25_index(dataset, collection_path, index_path):
|
||||
title = dataset['title']
|
||||
text = dataset['text']
|
||||
json_list = []
|
||||
for i in range(len(title)):
|
||||
json_dict = {'id': i, 'contents': title[i] + ' -- ' + text[i]}
|
||||
json_list.append(json_dict)
|
||||
|
||||
with open(os.path.join(collection_path, 'documents.json'), 'w') as f:
|
||||
json.dump(json_list, f)
|
||||
|
||||
command = f"python -u -m pyserini.index.lucene --collection JsonCollection --input {collection_path} --index {index_path} --generator DefaultLuceneDocumentGenerator --threads 8 --storePositions --storeDocvectors --storeRaw"
|
||||
result = subprocess.run(command, capture_output=True, text=True, shell=True)
|
||||
if result.returncode == 0:
|
||||
output = result.stdout
|
||||
print("execute successful!")
|
||||
print(output)
|
||||
else:
|
||||
print("execute false!")
|
||||
print(result.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = HfArgumentParser((ModelArguments, DataArguments))
|
||||
model_args, data_args = parser.parse_args_into_dataclasses()
|
||||
dataset_path = os.path.join(data_args.data_path, 'dataset')
|
||||
collection_path = os.path.join(data_args.data_path, 'collection')
|
||||
index_path = os.path.join(data_args.data_path, 'index')
|
||||
emb_path = os.path.join(data_args.data_path, 'emb')
|
||||
os.makedirs(dataset_path, exist_ok=True)
|
||||
os.makedirs(collection_path, exist_ok=True)
|
||||
os.makedirs(index_path, exist_ok=True)
|
||||
os.makedirs(emb_path, exist_ok=True)
|
||||
dataset = load_dataset(f"Cohere/wikipedia-22-12", 'zh', split='train')
|
||||
dataset.save_to_disk(dataset_path)
|
||||
build_bm25_index(dataset, collection_path, index_path)
|
||||
inference(os.path.join(collection_path, 'documents.json'),
|
||||
os.path.join(emb_path, 'data.npy'),
|
||||
model_args.model_name_or_path)
|
||||
@@ -0,0 +1,66 @@
|
||||
# Q&A Example
|
||||
|
||||
Vector Database can help LLMs to access external knowledge.
|
||||
You can load baai-general-embedding as the encoder to generate the vectors.
|
||||
Here a example to build a bot which can answer your question using the knowledge in chinese wikipedia.
|
||||
|
||||
Here's a description of the Q&A dialogue scenario using flag embedding and a large language model:
|
||||
|
||||
1. **Data Preprocessing and Indexing:**
|
||||
- Download a Chinese wikipedia dataset.
|
||||
- Encode the Chinese wikipedia text using flag embedding.
|
||||
- Build an index using BM25.
|
||||
2. **Query Enhancement with Large Language Model (LLM):**
|
||||
- Utilize a Large Language Model (LLM) to enhance and enrich the original user query based on the chat history.
|
||||
- The LLM can perform tasks such as text completion and paraphrasing to make the query more robust and comprehensive.
|
||||
3. **Document Retrieval:**
|
||||
- Employ BM25 to retrieve the top-n documents from the locally stored Chinese wiki dataset based on the newly enhanced query.
|
||||
4. **Embedding Retrieval:**
|
||||
- Perform an embedding retrieval on the top-n retrieved documents using brute force search to get top-k documents.
|
||||
5. **Answer Retrieval with Language Model (LLM):**
|
||||
- Present the question, the top-k retrieved documents, and chat history to the Large Language Model (LLM).
|
||||
- The LLM can utilize its understanding of language and context to provide accurate and comprehensive answers to the user's question.
|
||||
|
||||
By following these steps, the Q&A system can leverage flag embedding, BM25 indexing, and a Large Language Model to improve the accuracy and intelligence of the system. The integration of these techniques can create a more sophisticated and reliable Q&A system for users, providing them with comprehensive information to effectively answer their questions.
|
||||
|
||||
### Installation
|
||||
|
||||
```shell
|
||||
sudo apt install default-jdk
|
||||
pip install -r requirements.txt
|
||||
conda install -c anaconda openjdk
|
||||
```
|
||||
|
||||
### Prepare Data
|
||||
|
||||
```shell
|
||||
python pre_process.py --data_path ./data
|
||||
```
|
||||
|
||||
This script will download the dataset (Chinese wikipedia), building BM25 index, inference embedding, and then save them to `data_path`.
|
||||
|
||||
## Q&A usage
|
||||
|
||||
### Run Directly
|
||||
|
||||
```shell
|
||||
export OPENAI_API_KEY=...
|
||||
python run.py --data_path ./data
|
||||
```
|
||||
|
||||
This script will build a Q&A dialogue scenario.
|
||||
|
||||
### Quick Start
|
||||
|
||||
```python
|
||||
# encoding=gbk
|
||||
from tool import LocalDatasetLoader, BMVectorIndex, Agent
|
||||
loader = LocalDatasetLoader(data_path="./data/dataset",
|
||||
embedding_path="./data/emb/data.npy")
|
||||
index = BMVectorIndex(model_path="BAAI/bge-large-zh",
|
||||
bm_index_path="./data/index",
|
||||
data_loader=loader)
|
||||
agent = Agent(index)
|
||||
question = "上次有人登月是什么时候"
|
||||
agent.Answer(question, RANKING=1000, TOP_N=5, verbose=False)
|
||||
```
|
||||
@@ -0,0 +1,12 @@
|
||||
datasets==2.14.0
|
||||
faiss-gpu==1.7.2
|
||||
langchain==0.0.244
|
||||
numpy==1.23.3
|
||||
pyserini==0.21.0
|
||||
tiktoken==0.4.0
|
||||
torch==2.0.1
|
||||
torch_geometric==2.3.1
|
||||
tqdm==4.65.0
|
||||
transformers==4.30.2
|
||||
openai==0.27.4
|
||||
urllib3==1.25.11
|
||||
@@ -0,0 +1,26 @@
|
||||
# encoding=gbk
|
||||
import os
|
||||
import sys
|
||||
|
||||
sys.path.append('../')
|
||||
|
||||
from transformers import HfArgumentParser
|
||||
|
||||
from search_demo.tool import LocalDatasetLoader, BMVectorIndex, Agent
|
||||
from search_demo.arguments import ModelArguments, DataArguments
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = HfArgumentParser((ModelArguments, DataArguments))
|
||||
model_args, data_args = parser.parse_args_into_dataclasses()
|
||||
loader = LocalDatasetLoader(data_path=os.path.join(data_args.data_path, 'dataset'),
|
||||
embedding_path=os.path.join(data_args.data_path, 'emb/data.npy'))
|
||||
index = BMVectorIndex(model_path=model_args.model_name_or_path,
|
||||
bm_index_path=os.path.join(data_args.data_path, 'index'),
|
||||
data_loader=loader)
|
||||
agent = Agent(index)
|
||||
while True:
|
||||
question = input("ÎÊ£º").strip()
|
||||
if question != '':
|
||||
agent.answer(question, RANKING=1000, TOP_N=5, verbose=True)
|
||||
else:
|
||||
break
|
||||
@@ -0,0 +1,130 @@
|
||||
# encoding=gbk
|
||||
import faiss
|
||||
import numpy as np
|
||||
import tiktoken
|
||||
import torch
|
||||
from datasets import load_from_disk
|
||||
from langchain import PromptTemplate, LLMChain
|
||||
from langchain.chat_models import ChatOpenAI
|
||||
from pyserini.search.lucene import LuceneSearcher
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
|
||||
|
||||
class LocalDatasetLoader:
|
||||
data_path: str = ""
|
||||
doc_emb: np.ndarray = None
|
||||
|
||||
def __init__(self,
|
||||
data_path,
|
||||
embedding_path):
|
||||
dataset = load_from_disk(data_path)
|
||||
title = dataset['title']
|
||||
text = dataset['text']
|
||||
self.data = [title[i] + ' -- ' + text[i] for i in range(len(title))]
|
||||
self.doc_emb = np.load(embedding_path)
|
||||
|
||||
|
||||
class QueryGenerator:
|
||||
def __init__(self):
|
||||
prompt_template = """请基于历史，对新问题进行修改，如果新问题与历史相关，你必须结合语境将代词替换为相应的指代内容，让它的提问更加明确；否则保留原始的新问题。只修改新问题，不作任何回答：\n历史：{history}\n新问题：{question}\n修改后的新问题："""
|
||||
llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo-0613")
|
||||
prompt = PromptTemplate(
|
||||
input_variables=["history", "question"],
|
||||
template=prompt_template,
|
||||
)
|
||||
|
||||
self.llm_chain = LLMChain(llm=llm, prompt=prompt)
|
||||
|
||||
def run(self, history, question):
|
||||
return self.llm_chain.predict(history=history, question=question).strip()
|
||||
|
||||
|
||||
class AnswerGenerator:
|
||||
def __init__(self):
|
||||
prompt_template = """请基于参考，回答问题，不需要标注任何引用：\n历史：{history}\n问题：{question}\n参考：{references}\n答案："""
|
||||
llm = ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo-0613")
|
||||
prompt = PromptTemplate(
|
||||
input_variables=["history", "question", "references"],
|
||||
template=prompt_template,
|
||||
)
|
||||
|
||||
self.llm_chain = LLMChain(llm=llm, prompt=prompt)
|
||||
|
||||
def run(self, history, question, references):
|
||||
return self.llm_chain.predict(history=history, question=question, references=references).strip()
|
||||
|
||||
|
||||
class BMVectorIndex:
|
||||
def __init__(self,
|
||||
model_path,
|
||||
bm_index_path,
|
||||
data_loader):
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(model_path)
|
||||
self.model = AutoModel.from_pretrained(model_path)
|
||||
self.bm_searcher = LuceneSearcher(bm_index_path)
|
||||
self.loader = data_loader
|
||||
|
||||
if '-en' in model_path:
|
||||
raise NotImplementedError("only support chinese currently")
|
||||
|
||||
if 'noinstruct' in model_path:
|
||||
self.instruction = None
|
||||
else:
|
||||
self.instruction = "为这个句子生成表示以用于检索相关文章："
|
||||
|
||||
def search_for_doc(self, query: str, RANKING: int = 1000, TOP_N: int = 5):
|
||||
hits = self.bm_searcher.search(query, RANKING)
|
||||
ids = [int(e.docid) for e in hits]
|
||||
use_docs = self.loader.doc_emb[ids]
|
||||
|
||||
if self.instruction is not None:
|
||||
query = self.instruction + query
|
||||
encoded_input = self.tokenizer(query, max_length=512, padding=True, truncation=True, return_tensors='pt')
|
||||
with torch.no_grad():
|
||||
model_output = self.model(**encoded_input)
|
||||
query_emb = model_output.last_hidden_state[:, 0]
|
||||
query_emb = torch.nn.functional.normalize(query_emb, p=2, dim=-1).detach().cpu().numpy()
|
||||
|
||||
temp_index = faiss.IndexFlatIP(use_docs[0].shape[0])
|
||||
temp_index.add(use_docs)
|
||||
|
||||
_, I = temp_index.search(query_emb, TOP_N)
|
||||
|
||||
return "\n".join([self.loader.data[ids[I[0][i]]] for i in range(TOP_N)])
|
||||
|
||||
|
||||
class Agent:
|
||||
def __init__(self, index):
|
||||
self.memory = ""
|
||||
self.index = index
|
||||
self.query_generator = QueryGenerator()
|
||||
self.answer_generator = AnswerGenerator()
|
||||
|
||||
def empty_memory(self):
|
||||
self.memory = ""
|
||||
|
||||
def update_memory(self, question, answer):
|
||||
self.memory += f"问：{question}\n答：{answer}\n"
|
||||
encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
|
||||
while len(encoding.encode(self.memory)) > 3500:
|
||||
pos = self.memory[2:].rfind("问")
|
||||
self.memory = self.memory[pos:]
|
||||
|
||||
def generate_query(self, question):
|
||||
if self.memory == "":
|
||||
return question
|
||||
else:
|
||||
return self.query_generator.run(self.memory, question)
|
||||
|
||||
def generate_answer(self, query, references):
|
||||
return self.answer_generator.run(self.memory, query, references)
|
||||
|
||||
def answer(self, question, RANKING=1000, TOP_N=5, verbose=True):
|
||||
query = self.generate_query(question)
|
||||
references = self.index.search_for_doc(query, RANKING=RANKING, TOP_N=TOP_N)
|
||||
answer = self.generate_answer(question, references)
|
||||
self.update_memory(question, answer)
|
||||
if verbose:
|
||||
print('\033[96m' + "查：" + query + '\033[0m')
|
||||
print('\033[96m' + "参：" + references + '\033[0m')
|
||||
print("答：" + answer)
|
||||
@@ -0,0 +1,76 @@
|
||||
# Unified Finetune
|
||||
|
||||
In this example, we show how to perform unified fine-tuning based on [`BAAI/bge-m3`](https://huggingface.co/BAAI/bge-m3) with your data.
|
||||
|
||||
## 1. Installation
|
||||
|
||||
```
|
||||
git clone https://github.com/FlagOpen/FlagEmbedding.git
|
||||
cd FlagEmbedding/research/BGE_M3
|
||||
```
|
||||
|
||||
|
||||
## 2. Data format
|
||||
|
||||
Training data should be a jsonl file, where each line is a dict like this:
|
||||
|
||||
```
|
||||
{"query": str, "pos": List[str], "neg":List[str]}
|
||||
```
|
||||
|
||||
`query` is the query, and `pos` is a list of positive texts, `neg` is a list of negative texts.
|
||||
|
||||
If you want to use knowledge distillation, each line of your jsonl file should be like this:
|
||||
|
||||
```
|
||||
{"query": str, "pos": List[str], "neg":List[str], "pos_scores": List[float], "neg_scores": List[float]}
|
||||
```
|
||||
|
||||
`pos_scores` is a list of positive scores, where `pos_scores[i]` is the score between `query` and `pos[i]` from the teacher model. `neg_scores` is a list of negative scores, where `neg_scores[i]` is the score between `query` and `neg[i]` from the teacher model.
|
||||
|
||||
See [toy_train_data](./toy_train_data) for an example of training data.
|
||||
|
||||
|
||||
## 3. Train
|
||||
|
||||
> **Note**: If you only want to fine-tune the dense embedding of `BAAI/bge-m3`, you can refer to [here](https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/finetune/embedder#1-standard-model).
|
||||
|
||||
Here is an simple example of how to perform unified fine-tuning (dense embedding, sparse embedding and colbert) based on `BAAI/bge-m3`:
|
||||
|
||||
```bash
|
||||
torchrun --nproc_per_node {number of gpus} \
|
||||
-m run \
|
||||
--output_dir {path to save model} \
|
||||
--model_name_or_path BAAI/bge-m3 \
|
||||
--train_data ./toy_train_data \
|
||||
--learning_rate 1e-5 \
|
||||
--fp16 \
|
||||
--num_train_epochs 5 \
|
||||
--per_device_train_batch_size {large batch size; set 1 for toy data} \
|
||||
--dataloader_drop_last True \
|
||||
--normlized True \
|
||||
--temperature 0.02 \
|
||||
--query_max_len 64 \
|
||||
--passage_max_len 256 \
|
||||
--train_group_size 2 \
|
||||
--negatives_cross_device \
|
||||
--logging_steps 10 \
|
||||
--same_task_within_batch True \
|
||||
--unified_finetuning True \
|
||||
--use_self_distill True
|
||||
```
|
||||
|
||||
You can also refer to [this script](./unified_finetune_bge-m3_exmaple.sh) for more details. In this script, we use `deepspeed` to perform distributed training. Learn more about `deepspeed` at https://www.deepspeed.ai/getting-started/. Note that there are some important parameters to be modified in this script:
|
||||
|
||||
- `HOST_FILE_CONTENT`: Machines and GPUs for training. If you want to use multiple machines for training, please refer to https://www.deepspeed.ai/getting-started/#resource-configuration-multi-node (note that you should configure `pdsh` and `ssh` properly).
|
||||
- `DS_CONFIG_FILE`: Path of deepspeed config file. [Here](https://github.com/FlagOpen/FlagEmbedding/blob/master/examples/finetune/ds_stage0.json) is an example of `ds_config.json`.
|
||||
- `DATA_PATH`: One or more paths of training data. **Each path must be a directory containing one or more jsonl files**.
|
||||
- `DEFAULT_BATCH_SIZE`: Default batch size for training. If you use efficient batching strategy, which means you have split your data to different parts by sequence length, then the batch size for each part will be decided by the `get_file_batch_size()` function in [`BGE_M3/data.py`](../../BGE_M3/data.py). Before starting training, you should set the corresponding batch size for each part in this function according to the GPU memory of your machines. `DEFAULT_BATCH_SIZE` will be used for the part whose sequence length is not in the `get_file_batch_size()` function.
|
||||
- `EPOCHS`: Number of training epochs.
|
||||
- `LEARNING_RATE`: The initial learning rate.
|
||||
- `SAVE_PATH`: Path of saving finetuned model.
|
||||
|
||||
You should set these parameters appropriately.
|
||||
|
||||
|
||||
For more detailed arguments setting, please refer to [`BGE_M3/arguments.py`](../../BGE_M3/arguments.py).
|
||||
@@ -0,0 +1,10 @@
|
||||
{"query": "Five women walk along a beach wearing flip-flops.", "pos": ["Some women with flip-flops on, are walking along the beach"], "neg": ["The 4 women are sitting on the beach.", "There was a reform in 1996.", "She's not going to court to clear her record.", "The man is talking about hawaii.", "A woman is standing outside.", "The battle was over. ", "A group of people plays volleyball."]}
|
||||
{"query": "A woman standing on a high cliff on one leg looking over a river.", "pos": ["A woman is standing on a cliff."], "neg": ["A woman sits on a chair.", "George Bush told the Republicans there was no way he would let them even consider this foolish idea, against his top advisors advice.", "The family was falling apart.", "no one showed up to the meeting", "A boy is sitting outside playing in the sand.", "Ended as soon as I received the wire.", "A child is reading in her bedroom."]}
|
||||
{"query": "Two woman are playing instruments; one a clarinet, the other a violin.", "pos": ["Some people are playing a tune."], "neg": ["Two women are playing a guitar and drums.", "A man is skiing down a mountain.", "The fatal dose was not taken when the murderer thought it would be.", "Person on bike", "The girl is standing, leaning against the archway.", "A group of women watch soap operas.", "No matter how old people get they never forget. "]}
|
||||
{"query": "A girl with a blue tank top sitting watching three dogs.", "pos": ["A girl is wearing blue."], "neg": ["A girl is with three cats.", "The people are watching a funeral procession.", "The child is wearing black.", "Financing is an issue for us in public schools.", "Kids at a pool.", "It is calming to be assaulted.", "I face a serious problem at eighteen years old. "]}
|
||||
{"query": "A yellow dog running along a forest path.", "pos": ["a dog is running"], "neg": ["a cat is running", "Steele did not keep her original story.", "The rule discourages people to pay their child support.", "A man in a vest sits in a car.", "Person in black clothing, with white bandanna and sunglasses waits at a bus stop.", "Neither the Globe or Mail had comments on the current state of Canada's road system. ", "The Spring Creek facility is old and outdated."]}
|
||||
{"query": "It sets out essential activities in each phase along with critical factors related to those activities.", "pos": ["Critical factors for essential activities are set out."], "neg": ["It lays out critical activities but makes no provision for critical factors related to those activities.", "People are assembled in protest.", "The state would prefer for you to do that.", "A girl sits beside a boy.", "Two males are performing.", "Nobody is jumping", "Conrad was being plotted against, to be hit on the head."]}
|
||||
{"query": "A man giving a speech in a restaurant.", "pos": ["A person gives a speech."], "neg": ["The man sits at the table and eats food.", "This is definitely not an endorsement.", "They sold their home because they were retiring and not because of the loan.", "The seal of Missouri is perfect.", "Someone is raising their hand.", "An athlete is competing in the 1500 meter swimming competition.", "Two men watching a magic show."]}
|
||||
{"query": "Indians having a gathering with coats and food and drinks.", "pos": ["A group of Indians are having a gathering with food and drinks"], "neg": ["A group of Indians are having a funeral", "It is only staged on Winter afternoons in Palma's large bullring.", "Right information can empower the legal service practices and the justice system. ", "Meanwhile, the mainland was empty of population.", "Two children is sleeping.", "a fisherman is trying to catch a monkey", "the people are in a train"]}
|
||||
{"query": "A woman with violet hair rides her bicycle outside.", "pos": ["A woman is riding her bike."], "neg": ["A woman is jogging in the park.", "The street was lined with white-painted houses.", "A group watches a movie inside.", "man at picnics cut steak", "Several chefs are sitting down and talking about food.", "The Commission notes that no significant alternatives were considered.", "We ran out of firewood and had to use pine needles for the fire."]}
|
||||
{"query": "A man pulls two women down a city street in a rickshaw.", "pos": ["A man is in a city."], "neg": ["A man is a pilot of an airplane.", "It is boring and mundane.", "The morning sunlight was shining brightly and it was warm. ", "Two people jumped off the dock.", "People watching a spaceship launch.", "Mother Teresa is an easy choice.", "It's worth being able to go at a pace you prefer."]}
|
||||
@@ -0,0 +1,10 @@
|
||||
{"query": "Five women walk along a beach wearing flip-flops.", "pos": ["Some women with flip-flops on, are walking along the beach"], "neg": ["The 4 women are sitting on the beach.", "There was a reform in 1996.", "She's not going to court to clear her record.", "The man is talking about hawaii.", "A woman is standing outside.", "The battle was over. ", "A group of people plays volleyball."]}
|
||||
{"query": "A woman standing on a high cliff on one leg looking over a river.", "pos": ["A woman is standing on a cliff."], "neg": ["A woman sits on a chair.", "George Bush told the Republicans there was no way he would let them even consider this foolish idea, against his top advisors advice.", "The family was falling apart.", "no one showed up to the meeting", "A boy is sitting outside playing in the sand.", "Ended as soon as I received the wire.", "A child is reading in her bedroom."]}
|
||||
{"query": "Two woman are playing instruments; one a clarinet, the other a violin.", "pos": ["Some people are playing a tune."], "neg": ["Two women are playing a guitar and drums.", "A man is skiing down a mountain.", "The fatal dose was not taken when the murderer thought it would be.", "Person on bike", "The girl is standing, leaning against the archway.", "A group of women watch soap operas.", "No matter how old people get they never forget. "]}
|
||||
{"query": "A girl with a blue tank top sitting watching three dogs.", "pos": ["A girl is wearing blue."], "neg": ["A girl is with three cats.", "The people are watching a funeral procession.", "The child is wearing black.", "Financing is an issue for us in public schools.", "Kids at a pool.", "It is calming to be assaulted.", "I face a serious problem at eighteen years old. "]}
|
||||
{"query": "A yellow dog running along a forest path.", "pos": ["a dog is running"], "neg": ["a cat is running", "Steele did not keep her original story.", "The rule discourages people to pay their child support.", "A man in a vest sits in a car.", "Person in black clothing, with white bandanna and sunglasses waits at a bus stop.", "Neither the Globe or Mail had comments on the current state of Canada's road system. ", "The Spring Creek facility is old and outdated."]}
|
||||
{"query": "It sets out essential activities in each phase along with critical factors related to those activities.", "pos": ["Critical factors for essential activities are set out."], "neg": ["It lays out critical activities but makes no provision for critical factors related to those activities.", "People are assembled in protest.", "The state would prefer for you to do that.", "A girl sits beside a boy.", "Two males are performing.", "Nobody is jumping", "Conrad was being plotted against, to be hit on the head."]}
|
||||
{"query": "A man giving a speech in a restaurant.", "pos": ["A person gives a speech."], "neg": ["The man sits at the table and eats food.", "This is definitely not an endorsement.", "They sold their home because they were retiring and not because of the loan.", "The seal of Missouri is perfect.", "Someone is raising their hand.", "An athlete is competing in the 1500 meter swimming competition.", "Two men watching a magic show."]}
|
||||
{"query": "Indians having a gathering with coats and food and drinks.", "pos": ["A group of Indians are having a gathering with food and drinks"], "neg": ["A group of Indians are having a funeral", "It is only staged on Winter afternoons in Palma's large bullring.", "Right information can empower the legal service practices and the justice system. ", "Meanwhile, the mainland was empty of population.", "Two children is sleeping.", "a fisherman is trying to catch a monkey", "the people are in a train"]}
|
||||
{"query": "A woman with violet hair rides her bicycle outside.", "pos": ["A woman is riding her bike."], "neg": ["A woman is jogging in the park.", "The street was lined with white-painted houses.", "A group watches a movie inside.", "man at picnics cut steak", "Several chefs are sitting down and talking about food.", "The Commission notes that no significant alternatives were considered.", "We ran out of firewood and had to use pine needles for the fire."]}
|
||||
{"query": "A man pulls two women down a city street in a rickshaw.", "pos": ["A man is in a city."], "neg": ["A man is a pilot of an airplane.", "It is boring and mundane.", "The morning sunlight was shining brightly and it was warm. ", "Two people jumped off the dock.", "People watching a spaceship launch.", "Mother Teresa is an easy choice.", "It's worth being able to go at a pace you prefer."]}
|
||||
@@ -0,0 +1,90 @@
|
||||
#!/bin/bash
|
||||
# Set root path
|
||||
ROOT=/home
|
||||
|
||||
# Set training machines
|
||||
# For more details, refer to https://www.deepspeed.ai/getting-started/#resource-configuration-multi-node
|
||||
HOST_FILE_CONTENT="\
|
||||
localhost slots=8\n\
|
||||
"
|
||||
HOST_FILE=hostfile
|
||||
printf "$HOST_FILE_CONTENT" > $HOST_FILE
|
||||
|
||||
DISTRIBUTED_ARGS="--hostfile $HOST_FILE"
|
||||
|
||||
export LAUNCHER="deepspeed \
|
||||
$DISTRIBUTED_ARGS \
|
||||
"
|
||||
# Set cache directory
|
||||
CACHE_PATH=$ROOT/datasets/.cache
|
||||
|
||||
# Set path of deepspeed config file
|
||||
# For more details, refer to https://huggingface.co/docs/transformers/main_classes/deepspeed#zero
|
||||
DS_CONFIG_FILE=$ROOT/train/ds_config.json
|
||||
|
||||
# Set group size of training
|
||||
GROUP_SIZE=2
|
||||
|
||||
# Set paths of training data. Every path **must be a directory path**.
|
||||
DATA_PATH="
|
||||
$ROOT/datasets/toy_train_data \
|
||||
"
|
||||
|
||||
# Set default batch size for training.
|
||||
# If you want to use effient batching strategy, you should use the script `split_data_by_length.py` to split your data by sequence length firstly. Then the batch size for every batch will depend on its sequence length range, such as len-0-500: 48, len-500-1000: 32, etc., which are defined in `get_file_batch_size()` in `BGE_M3/data.py`.
|
||||
DEFAULT_BATCH_SIZE=1
|
||||
|
||||
# Set number of training epochs.
|
||||
# For more details, refer to https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments
|
||||
EPOCHS=5
|
||||
MAX_STEPS=-1
|
||||
|
||||
# Set base model and save path
|
||||
BASE_MODEL=BAAI/bge-m3
|
||||
|
||||
SAVE_PATH=$ROOT/models/bge-m3_finetuned
|
||||
mkdir -p $SAVE_PATH
|
||||
|
||||
# Set learning rate
|
||||
LEARNING_RATE=5e-6
|
||||
|
||||
full_options="
|
||||
--knowledge_distillation True \
|
||||
--output_dir $SAVE_PATH \
|
||||
--model_name_or_path $BASE_MODEL \
|
||||
--normlized True \
|
||||
--temperature 0.02 \
|
||||
--do_train \
|
||||
--train_data $DATA_PATH \
|
||||
--cache_path $CACHE_PATH \
|
||||
--per_device_train_batch_size $DEFAULT_BATCH_SIZE \
|
||||
--query_max_len 512 \
|
||||
--passage_max_len 8192 \
|
||||
--small_threshold 200 \
|
||||
--drop_threshold 200 \
|
||||
--fp16 \
|
||||
--save_steps 1500 \
|
||||
--train_group_size $GROUP_SIZE \
|
||||
--learning_rate $LEARNING_RATE \
|
||||
--num_train_epochs $EPOCHS \
|
||||
--max_steps $MAX_STEPS \
|
||||
--negatives_cross_device False \
|
||||
--logging_steps 10 \
|
||||
--warmup_ratio 0.1 \
|
||||
--weight_decay 0.01 \
|
||||
--overwrite_output_dir True \
|
||||
--gradient_checkpointing \
|
||||
--sentence_pooling_method cls \
|
||||
--same_task_within_batch True \
|
||||
--shuffle_ratio 0.002 \
|
||||
--enable_sub_batch True \
|
||||
--deepspeed ${DS_CONFIG_FILE} \
|
||||
--unified_finetuning True \
|
||||
--use_self_distill True
|
||||
"
|
||||
|
||||
run_cmd="$LAUNCHER --module FlagEmbedding.BGE_M3.run ${full_options}"
|
||||
echo ${run_cmd}
|
||||
eval ${run_cmd} 2>&1 | tee $SAVE_PATH/output.log
|
||||
|
||||
set +x
|
||||
Reference in New Issue
Block a user