chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,200 @@
|
||||
<div align="center">
|
||||
<h1> Llama2Vec: Unsupervised Adaptation of Large Language Models for Dense Retrieval (LLARA) [<a href="https://arxiv.org/abs/2312.15503">paper</a>]</h1>
|
||||
</div>
|
||||
|
||||
Llama2Vec consists of two pretext tasks:
|
||||
- **EBAE** (Embedding-Based Auto-Encoding)
|
||||
- **EBAR** (Embedding-Based Auto-Regression)
|
||||
|
||||
The LLM is prompted to **reconstruct the input sentence** and **predict the next sentence** based on its text embeddings.
|
||||
|
||||
It is known for the following features:
|
||||
- simple
|
||||
- lightweight
|
||||
- highly effective
|
||||
|
||||
## Environment
|
||||
```bash
|
||||
conda create llara python=3.10
|
||||
|
||||
conda activate llara
|
||||
|
||||
# You may need to adjust the cuda version
|
||||
conda install pytorch pytorch-cuda=12.1 -c pytorch -c nvidia
|
||||
pip install transformers==4.41.0 deepspeed accelerate datasets peft pandas
|
||||
pip install flash-attn --no-build-isolation
|
||||
```
|
||||
|
||||
## Model List
|
||||
|
||||
| Model | Introduction |
|
||||
| ------------------------------------------------------------ | ------------------------------------------------------------ |
|
||||
| [BAAI/LLARA-pretrain](https://huggingface.co/BAAI/LLARA-pretrain) | LLARA that has undergone unsupervised adaptation on Wikipedia |
|
||||
| [BAAI/LLARA-passage](https://huggingface.co/BAAI/LLARA-passage) | The LLARA-pretrain model fine-tuned on MS MARCO passage (the hard negatives come from dense retriever) |
|
||||
| [BAAI/LLARA-document](https://huggingface.co/BAAI/LLARA-document) | The LLARA-pretrain model fine-tuned on MS MARCO document |
|
||||
| [BAAI/LLARA-beir](https://huggingface.co/BAAI/LLARA-beir) | The LLARA-pretrain model fine-tuned on MS MARCO passage (the hard negatives come from BM25) |
|
||||
|
||||
## Usage
|
||||
|
||||
```python
|
||||
import torch
|
||||
from transformers import AutoModel, AutoTokenizer, LlamaModel
|
||||
|
||||
def get_query_inputs(queries, tokenizer, max_length=512):
|
||||
prefix = '"'
|
||||
suffix = '", predict the following passage within eight words: <s9><s10><s11><s12><s13><s14><s15><s16>'
|
||||
prefix_ids = tokenizer(prefix, return_tensors=None)['input_ids']
|
||||
suffix_ids = tokenizer(suffix, return_tensors=None)['input_ids'][1:]
|
||||
queries_inputs = []
|
||||
for query in queries:
|
||||
inputs = tokenizer(query,
|
||||
return_tensors=None,
|
||||
max_length=max_length,
|
||||
truncation=True,
|
||||
add_special_tokens=False)
|
||||
inputs['input_ids'] = prefix_ids + inputs['input_ids'] + suffix_ids
|
||||
inputs['attention_mask'] = [1] * len(inputs['input_ids'])
|
||||
queries_inputs.append(inputs)
|
||||
return tokenizer.pad(
|
||||
queries_inputs,
|
||||
padding=True,
|
||||
max_length=max_length,
|
||||
pad_to_multiple_of=8,
|
||||
return_tensors='pt',
|
||||
)
|
||||
|
||||
def get_passage_inputs(passages, tokenizer, max_length=512):
|
||||
prefix = '"'
|
||||
suffix = '", summarize the above passage within eight words: <s1><s2><s3><s4><s5><s6><s7><s8>'
|
||||
prefix_ids = tokenizer(prefix, return_tensors=None)['input_ids']
|
||||
suffix_ids = tokenizer(suffix, return_tensors=None)['input_ids'][1:]
|
||||
passages_inputs = []
|
||||
for passage in passages:
|
||||
inputs = tokenizer(passage,
|
||||
return_tensors=None,
|
||||
max_length=max_length,
|
||||
truncation=True,
|
||||
add_special_tokens=False)
|
||||
inputs['input_ids'] = prefix_ids + inputs['input_ids'] + suffix_ids
|
||||
inputs['attention_mask'] = [1] * len(inputs['input_ids'])
|
||||
passages_inputs.append(inputs)
|
||||
return tokenizer.pad(
|
||||
passages_inputs,
|
||||
padding=True,
|
||||
max_length=max_length,
|
||||
pad_to_multiple_of=8,
|
||||
return_tensors='pt',
|
||||
)
|
||||
|
||||
# Load the tokenizer and model
|
||||
tokenizer = AutoTokenizer.from_pretrained('BAAI/LLARA-passage')
|
||||
model = AutoModel.from_pretrained('BAAI/LLARA-passage')
|
||||
|
||||
# Define query and passage inputs
|
||||
query = "What is llama?"
|
||||
title = "Llama"
|
||||
passage = "The llama is a domesticated South American camelid, widely used as a meat and pack animal by Andean cultures since the pre-Columbian era."
|
||||
query_input = get_query_inputs([query], tokenizer)
|
||||
passage_input = get_passage_inputs([passage], tokenizer)
|
||||
|
||||
|
||||
with torch.no_grad():
|
||||
# compute query embedding
|
||||
query_outputs = model(**query_input, return_dict=True, output_hidden_states=True)
|
||||
query_embedding = query_outputs.hidden_states[-1][:, -8:, :]
|
||||
query_embedding = torch.mean(query_embedding, dim=1)
|
||||
query_embedding = torch.nn.functional.normalize(query_embedding, dim=-1)
|
||||
|
||||
# compute passage embedding
|
||||
passage_outputs = model(**passage_input, return_dict=True, output_hidden_states=True)
|
||||
passage_embeddings = passage_outputs.hidden_states[-1][:, -8:, :]
|
||||
passage_embeddings = torch.mean(passage_embeddings, dim=1)
|
||||
passage_embeddings = torch.nn.functional.normalize(passage_embeddings, dim=-1)
|
||||
|
||||
# compute similarity score
|
||||
score = query_embedding @ passage_embeddings.T
|
||||
print(score)
|
||||
|
||||
```
|
||||
|
||||
## Unsupervised Adaption (pretrain)
|
||||
1. You can get the complete data here: [cfli/pretrain_wiki](https://huggingface.co/datasets/cfli/pretrain_wiki)
|
||||
2. Here is an example for pretrain:
|
||||
```shell
|
||||
cd ./pretrain
|
||||
torchrun --nproc_per_node 8 \
|
||||
run.py \
|
||||
--output_dir ./output \
|
||||
--model_name_or_path meta-llama/Llama-2-7b-hf \
|
||||
--train_data ../data/pretrain/toy_pretrain_data.jsonl \
|
||||
--learning_rate 1e-5 \
|
||||
--num_train_epochs 1 \
|
||||
--per_device_train_batch_size 1 \
|
||||
--gradient_accumulation_steps 1 \
|
||||
--dataloader_drop_last True \
|
||||
--cutoff_len 128 \
|
||||
--logging_steps 1 \
|
||||
--save_steps 500 \
|
||||
--save_total_limit 20 \
|
||||
--gradient_checkpointing \
|
||||
--ddp_find_unused_parameters False \
|
||||
--use_flash_attn False \
|
||||
--deepspeed ../stage1.json \
|
||||
--warmup_ratio 0.1 \
|
||||
--remove_stop_words True \
|
||||
--use_lora False \
|
||||
--bf16 \
|
||||
--cache_dir ./LMs \
|
||||
--token ...
|
||||
```
|
||||
If you want to pretrain based on the complete data, please use hype-parameters in our paper.
|
||||
|
||||
## Fine-tune
|
||||
|
||||
Here is an example for fine-tune:
|
||||
```shell
|
||||
cd ./finetune
|
||||
torchrun --nproc_per_node 8 \
|
||||
run.py \
|
||||
--output_dir ./output \
|
||||
--model_name_or_path BAAI/LLARA-pretrain \
|
||||
--train_data ../data/finetune/toy_finetune_data.jsonl \
|
||||
--learning_rate 3e-4 \
|
||||
--num_train_epochs 1 \
|
||||
--per_device_train_batch_size 1 \
|
||||
--dataloader_drop_last True \
|
||||
--normlized True \
|
||||
--temperature 0.01 \
|
||||
--query_max_len 64 \
|
||||
--passage_max_len 160 \
|
||||
--train_group_size 16 \
|
||||
--logging_steps 10 \
|
||||
--save_steps 500 \
|
||||
--save_total_limit 3 \
|
||||
--ddp_find_unused_parameters False \
|
||||
--negatives_cross_device \
|
||||
--gradient_checkpointing \
|
||||
--deepspeed ../stage1.json \
|
||||
--warmup_ratio 0.1 \
|
||||
--fp16 \
|
||||
--cache_dir ./LMs \
|
||||
--token ...
|
||||
```
|
||||
|
||||
## Citation
|
||||
|
||||
If you find this repository useful, please give us a star ⭐.
|
||||
|
||||
To cite our work:
|
||||
|
||||
```
|
||||
@misc{li2023makinglargelanguagemodels,
|
||||
title={Making Large Language Models A Better Foundation For Dense Retrieval},
|
||||
author={Chaofan Li and Zheng Liu and Shitao Xiao and Yingxia Shao},
|
||||
year={2023},
|
||||
eprint={2312.15503},
|
||||
archivePrefix={arXiv},
|
||||
primaryClass={cs.CL},
|
||||
url={https://arxiv.org/abs/2312.15503},
|
||||
}
|
||||
```
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -0,0 +1,165 @@
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, List
|
||||
|
||||
from transformers import TrainingArguments
|
||||
|
||||
|
||||
def default_list() -> List[int]:
|
||||
return ['v_proj', 'q_proj', 'k_proj', 'gate_proj', 'down_proj', 'o_proj', 'up_proj']
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArguments:
|
||||
"""
|
||||
Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
|
||||
"""
|
||||
|
||||
model_name_or_path: str = field(
|
||||
metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
|
||||
)
|
||||
|
||||
peft_model_path: str = field(
|
||||
default=''
|
||||
)
|
||||
config_name: Optional[str] = field(
|
||||
default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
|
||||
)
|
||||
tokenizer_name: Optional[str] = field(
|
||||
default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
|
||||
)
|
||||
# cache_dir: Optional[str] = field(
|
||||
# default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
|
||||
# )
|
||||
use_lora: bool = field(
|
||||
default=True,
|
||||
metadata={"help": "If passed, will use LORA (low-rank parameter-efficient training) to train the model."}
|
||||
)
|
||||
lora_rank: int = field(
|
||||
default=64,
|
||||
metadata={"help": "The rank of lora."}
|
||||
)
|
||||
lora_alpha: float = field(
|
||||
default=16,
|
||||
metadata={"help": "The alpha parameter of lora."}
|
||||
)
|
||||
lora_dropout: float = field(
|
||||
default=0.1,
|
||||
metadata={"help": "The dropout rate of lora modules."}
|
||||
)
|
||||
target_modules: List[str] = field(
|
||||
default_factory=default_list
|
||||
)
|
||||
save_merged_lora_model: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "If passed, will merge the lora modules and save the entire model."}
|
||||
)
|
||||
use_flash_attn: bool = field(
|
||||
default=True,
|
||||
metadata={"help": "If passed, will use flash attention to train the model."}
|
||||
)
|
||||
use_slow_tokenizer: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "If passed, will use a slow tokenizer (not backed by the 🤗 Tokenizers library)."}
|
||||
)
|
||||
low_cpu_mem_usage: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "It is an option to create the model as an empty shell,"
|
||||
"then only materialize its parameters when the pretrained weights are loaded."
|
||||
"If passed, LLM loading time and RAM consumption will be benefited."}
|
||||
)
|
||||
token: str = field(
|
||||
default=""
|
||||
)
|
||||
cache_dir: str = field(
|
||||
default="./LMs"
|
||||
)
|
||||
from_peft: str = field(
|
||||
default=None
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataArguments:
|
||||
train_data: str = field(
|
||||
default='./toy_finetune_data.jsonl', metadata={"help": "Path to train data"}
|
||||
)
|
||||
train_group_size: int = field(default=8)
|
||||
|
||||
query_max_len: int = field(
|
||||
default=32,
|
||||
metadata={
|
||||
"help": "The maximum total input sequence length after tokenization for passage. Sequences longer "
|
||||
"than this will be truncated, sequences shorter will be padded."
|
||||
},
|
||||
)
|
||||
|
||||
passage_max_len: int = field(
|
||||
default=128,
|
||||
metadata={
|
||||
"help": "The maximum total input sequence length after tokenization for passage. Sequences longer "
|
||||
"than this will be truncated, sequences shorter will be padded."
|
||||
},
|
||||
)
|
||||
|
||||
max_example_num_per_dataset: int = field(
|
||||
default=100000000, metadata={"help": "the max number of examples for each dataset"}
|
||||
)
|
||||
|
||||
query_instruction_for_retrieval: str = field(
|
||||
default="query: ", metadata={"help": "query: "}
|
||||
)
|
||||
passage_instruction_for_retrieval: str = field(
|
||||
default="passage: ", metadata={"help": "passage: "}
|
||||
)
|
||||
|
||||
cache_path: str = field(
|
||||
default='./data_dir'
|
||||
)
|
||||
|
||||
load_from_disk: bool = field(
|
||||
default=False, metadata={"help": " whether load the data from disk"}
|
||||
)
|
||||
|
||||
load_disk_path: str = field(
|
||||
default=None, metadata={"help": " the path to load the data", "nargs": "+"}
|
||||
)
|
||||
|
||||
save_to_disk: bool = field(
|
||||
default=False, metadata={"help": " whether save the data to disk"}
|
||||
)
|
||||
|
||||
save_disk_path: str = field(
|
||||
default=None, metadata={"help": " the path to save the data"}
|
||||
)
|
||||
|
||||
num_shards: int = field(
|
||||
default=0, metadata={
|
||||
"help": "number of shards to write, prior than `save_max_shard_size`, default depends on `save_max_shard_size`"}
|
||||
)
|
||||
|
||||
save_max_shard_size: str = field(
|
||||
default="50GB", metadata={"help": "the max size of the shard"}
|
||||
)
|
||||
|
||||
exit_after_save: bool = field(
|
||||
default=False, metadata={"help": " whether exit after save the data"}
|
||||
)
|
||||
|
||||
shuffle_ratio: float = field(
|
||||
default=0.0, metadata={"help": "The ratio of shuffling the text"}
|
||||
)
|
||||
|
||||
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 RetrieverTrainingArguments(TrainingArguments):
|
||||
negatives_cross_device: bool = field(default=False, metadata={"help": "share negatives across devices"})
|
||||
temperature: Optional[float] = field(default=0.02)
|
||||
fix_position_embedding: bool = field(default=False, metadata={"help": "Freeze the parameters of position embeddings"})
|
||||
sentence_pooling_method: str = field(default='cls', metadata={"help": "the pooling method, should be cls or mean"})
|
||||
normlized: bool = field(default=True)
|
||||
sub_batch_size: int = field(default=None)
|
||||
cache_chunk_size: int = field(default=-1, metadata={"help": "用于缓存每一步的执行."})
|
||||
@@ -0,0 +1,161 @@
|
||||
import sys
|
||||
|
||||
import math
|
||||
import os.path
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Tuple
|
||||
|
||||
import datasets
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from transformers import DataCollatorWithPadding, DataCollatorForSeq2Seq
|
||||
from transformers import PreTrainedTokenizer, BatchEncoding
|
||||
|
||||
from arguments import DataArguments
|
||||
|
||||
|
||||
class TrainDatasetForEmbedding(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
args: DataArguments,
|
||||
tokenizer: PreTrainedTokenizer
|
||||
):
|
||||
if os.path.isdir(args.train_data):
|
||||
train_datasets = []
|
||||
for file in os.listdir(args.train_data):
|
||||
temp_dataset = datasets.load_dataset('json', data_files=os.path.join(args.train_data, file),
|
||||
split='train')
|
||||
if len(temp_dataset) > args.max_example_num_per_dataset:
|
||||
temp_dataset = temp_dataset.select(
|
||||
random.sample(list(range(len(temp_dataset))), args.max_example_num_per_dataset))
|
||||
train_datasets.append(temp_dataset)
|
||||
self.dataset = datasets.concatenate_datasets(train_datasets)
|
||||
else:
|
||||
self.dataset = datasets.load_dataset('json', data_files=args.train_data, split='train', cache_dir=args.cache_path)
|
||||
|
||||
self.tokenizer = tokenizer
|
||||
self.args = args
|
||||
self.total_len = len(self.dataset)
|
||||
|
||||
self.prefix = '"'
|
||||
self.prefix_ids = self.tokenizer(self.prefix, return_tensors=None)['input_ids']
|
||||
self.suffix_passage = '", summarize the above passage within eight words: <s1><s2><s3><s4><s5><s6><s7><s8>'
|
||||
self.suffix_passage_ids = self.tokenizer(self.suffix_passage, return_tensors=None, add_special_tokens=False)['input_ids']
|
||||
# self.suffix_query = '", predict the following passage within eight words: <s9><s10><s11><s12><s13><s14><s15><s16>'
|
||||
self.suffix_query = '", predict the following passage within eight words: <s9><s10><s11><s12><s13><s14><s15><s16>'
|
||||
self.suffix_query_ids = self.tokenizer(self.suffix_query, return_tensors=None, add_special_tokens=False)['input_ids']
|
||||
self.query_max_len = self.args.query_max_len - len(self.prefix_ids) - len(self.suffix_query_ids)
|
||||
self.passage_max_len = self.args.passage_max_len - len(self.prefix_ids) - len(self.suffix_passage_ids)
|
||||
|
||||
def __len__(self):
|
||||
# return self.total_len
|
||||
return self.total_len
|
||||
|
||||
def __getitem__(self, item) -> Tuple[BatchEncoding, List[BatchEncoding]]:
|
||||
query = self.dataset[item]['query']
|
||||
query_inputs = self.tokenizer(query,
|
||||
return_tensors=None,
|
||||
max_length=self.query_max_len,
|
||||
truncation=True,
|
||||
add_special_tokens=False)
|
||||
query_inputs['input_ids'] = self.prefix_ids + query_inputs['input_ids'] + self.suffix_query_ids
|
||||
query_inputs['attention_mask'] = [1] * len(query_inputs['input_ids'])
|
||||
|
||||
passages = []
|
||||
pos = random.choice(self.dataset[item]['pos'])
|
||||
passages.append(pos)
|
||||
|
||||
if len(self.dataset[item]['neg']) < self.args.train_group_size - 1:
|
||||
# print(len(self.dataset[item]['neg']))
|
||||
num = math.ceil((self.args.train_group_size - 1) / len(list(set(self.dataset[item]['neg']))))
|
||||
# negs = random.sample(list(set(self.dataset[item]['neg'])) * num, self.args.train_group_size - 1)
|
||||
negs = random.sample(self.dataset[item]['neg'] * num, self.args.train_group_size - 1)
|
||||
# negs = random.sample(self.dataset[item]['neg'], self.args.train_group_size - 1)
|
||||
else:
|
||||
negs = random.sample(self.dataset[item]['neg'], self.args.train_group_size - 1)
|
||||
# negs = random.sample(self.dataset[item]['neg'], self.args.train_group_size - 1)
|
||||
passages.extend(negs)
|
||||
|
||||
passages_inputs = []
|
||||
for passage in passages:
|
||||
passage_inputs = self.tokenizer(passage,
|
||||
return_tensors=None,
|
||||
max_length=self.passage_max_len,
|
||||
truncation=True,
|
||||
add_special_tokens=False)
|
||||
passage_inputs['input_ids'] = self.prefix_ids + passage_inputs['input_ids'] + self.suffix_passage_ids
|
||||
passage_inputs['attention_mask'] = [1] * len(passage_inputs['input_ids'])
|
||||
passages_inputs.append(passage_inputs)
|
||||
|
||||
return query_inputs, passages_inputs
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmbedCollator(DataCollatorForSeq2Seq):
|
||||
"""
|
||||
Wrapper that does conversion from List[Tuple[encode_qry, encode_psg]] to List[qry], List[psg]
|
||||
and pass batch separately to the actual collator.
|
||||
Abstract out data detail for the model.
|
||||
"""
|
||||
query_max_len: int = 32
|
||||
passage_max_len: int = 128
|
||||
sub_batch_size: int = -1
|
||||
|
||||
def __call__(self, features, return_tensors='pt'):
|
||||
if return_tensors is None:
|
||||
return_tensors = self.return_tensors
|
||||
|
||||
queries = []
|
||||
passages = []
|
||||
for e in features:
|
||||
queries.append(e[0])
|
||||
passages.extend(e[1])
|
||||
|
||||
if self.sub_batch_size is None or self.sub_batch_size <= 0:
|
||||
q_collated = self.tokenizer.pad(
|
||||
queries,
|
||||
padding=self.padding,
|
||||
max_length=self.query_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=return_tensors,
|
||||
)
|
||||
|
||||
d_collated = self.tokenizer.pad(
|
||||
passages,
|
||||
padding=self.padding,
|
||||
max_length=self.passage_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=return_tensors,
|
||||
)
|
||||
else:
|
||||
batch_size = self.sub_batch_size
|
||||
|
||||
q_collated = []
|
||||
for i in range(0, len(queries), batch_size):
|
||||
start = i
|
||||
end = min(len(queries), i + batch_size)
|
||||
sub_features = queries[start:end]
|
||||
q_collated.append(self.tokenizer.pad(
|
||||
sub_features,
|
||||
padding=self.padding,
|
||||
max_length=self.passage_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=return_tensors,
|
||||
))
|
||||
|
||||
d_collated = []
|
||||
for i in range(0, len(passages), batch_size):
|
||||
start = i
|
||||
end = min(len(passages), i + batch_size)
|
||||
sub_features = passages[start: end]
|
||||
d_collated.append(self.tokenizer.pad(
|
||||
sub_features,
|
||||
padding=self.padding,
|
||||
max_length=self.passage_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=return_tensors,
|
||||
))
|
||||
|
||||
return {"query": q_collated, "passage": d_collated}
|
||||
@@ -0,0 +1,65 @@
|
||||
import sys
|
||||
|
||||
import torch
|
||||
from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM, AutoModel
|
||||
from peft import LoraConfig, TaskType, get_peft_model, PeftModel
|
||||
|
||||
|
||||
def get_model(model_args):
|
||||
# if model_args.use_flash_attn:
|
||||
# from llama2_flash_attn_monkey_patch import replace_llama_attn_with_flash_attn
|
||||
# replace_llama_attn_with_flash_attn()
|
||||
|
||||
if model_args.config_name:
|
||||
config = AutoConfig.from_pretrained(model_args.config_name,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
)
|
||||
elif model_args.model_name_or_path:
|
||||
config = AutoConfig.from_pretrained(model_args.model_name_or_path,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
"You are instantiating a new config instance from scratch. This is not supported by this script."
|
||||
)
|
||||
config.use_cache = False
|
||||
|
||||
if model_args.model_name_or_path:
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_args.model_name_or_path,
|
||||
# load_in_8bit=True,
|
||||
# torch_dtype=torch.bfloat16,
|
||||
use_flash_attention_2=True if model_args.use_flash_attn else False,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
from_tf=bool(".ckpt" in model_args.model_name_or_path),
|
||||
config=config,
|
||||
# low_cpu_mem_usage=model_args.low_cpu_mem_usage,
|
||||
# device_map="auto",
|
||||
)
|
||||
else:
|
||||
print("Training new model from scratch")
|
||||
model = model_args.from_config(config)
|
||||
|
||||
if model_args.from_peft is not None:
|
||||
model = PeftModel.from_pretrained(model, model_args.from_peft, is_trainable=True)
|
||||
model.print_trainable_parameters()
|
||||
else:
|
||||
if model_args.use_lora:
|
||||
peft_config = LoraConfig(
|
||||
task_type=TaskType.FEATURE_EXTRACTION,
|
||||
inference_mode=False,
|
||||
r=model_args.lora_rank,
|
||||
target_modules=model_args.target_modules,
|
||||
lora_alpha=model_args.lora_alpha,
|
||||
lora_dropout=model_args.lora_dropout
|
||||
)
|
||||
model = get_peft_model(model, peft_config)
|
||||
# print(model.model.layers[0].self_attn.q_proj.weight.dtype)
|
||||
# print(model.model.layers[0].self_attn.q_proj.lora_A.default.weight.dtype)
|
||||
# sys.exit(0)
|
||||
model.print_trainable_parameters()
|
||||
|
||||
return model
|
||||
@@ -0,0 +1,165 @@
|
||||
import logging
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, Optional, List, Union
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from torch import nn, Tensor
|
||||
from tqdm import trange, tqdm
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
from transformers.file_utils import ModelOutput
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EncoderOutput(ModelOutput):
|
||||
q_reps: Optional[Tensor] = None
|
||||
p_reps: Optional[Tensor] = None
|
||||
loss: Optional[Tensor] = None
|
||||
scores: Optional[Tensor] = None
|
||||
|
||||
|
||||
class BiEncoderModel(nn.Module):
|
||||
TRANSFORMER_CLS = AutoModel
|
||||
|
||||
def __init__(self,
|
||||
model: AutoModel = None,
|
||||
tokenizer: AutoTokenizer = None,
|
||||
normlized: bool = False,
|
||||
negatives_cross_device: bool = False,
|
||||
temperature: float = 1.0,
|
||||
sub_batch_size: int = -1
|
||||
):
|
||||
super().__init__()
|
||||
self.model = model
|
||||
self.config = model.config
|
||||
self.tokenizer = tokenizer
|
||||
self.cross_entropy = nn.CrossEntropyLoss(reduction='mean')
|
||||
|
||||
self.normlized = normlized
|
||||
self.temperature = temperature
|
||||
if not normlized:
|
||||
self.temperature = 1.0
|
||||
logger.info("reset temperature = 1.0 due to using inner product to compute similarity")
|
||||
|
||||
self.negatives_cross_device = negatives_cross_device
|
||||
if self.negatives_cross_device:
|
||||
if not dist.is_initialized():
|
||||
raise ValueError('Distributed training has not been initialized for representation all gather.')
|
||||
# logger.info("Run in a single GPU, set negatives_cross_device=False")
|
||||
# self.negatives_cross_device = False
|
||||
# else:
|
||||
self.process_rank = dist.get_rank()
|
||||
self.world_size = dist.get_world_size()
|
||||
|
||||
self.sub_batch_size = sub_batch_size
|
||||
|
||||
def gradient_checkpointing_enable(self, **kwargs):
|
||||
self.model.gradient_checkpointing_enable(**kwargs)
|
||||
|
||||
def enable_input_require_grads(self, **kwargs):
|
||||
self.model.enable_input_require_grads(**kwargs)
|
||||
|
||||
def encode(self, features):
|
||||
# input('continue?')
|
||||
if features is None:
|
||||
return None
|
||||
if not isinstance(features, list):
|
||||
if self.sub_batch_size is not None and self.sub_batch_size > 0:
|
||||
all_p_reps = []
|
||||
for i in range(0, len(features['attention_mask']), self.sub_batch_size):
|
||||
end_inx = min(i + self.sub_batch_size, len(features['attention_mask']))
|
||||
sub_features = {}
|
||||
for k, v in features.items():
|
||||
sub_features[k] = v[i:end_inx]
|
||||
psg_out = self.model(**sub_features, return_dict=True, output_hidden_states=True)
|
||||
### modify
|
||||
p_reps = psg_out.hidden_states[-1][:, -8:, :]
|
||||
p_reps = torch.mean(p_reps, dim=1)
|
||||
all_p_reps.append(p_reps)
|
||||
all_p_reps = torch.cat(all_p_reps, 0).contiguous()
|
||||
if self.normlized:
|
||||
all_p_reps = torch.nn.functional.normalize(all_p_reps, dim=-1)
|
||||
return all_p_reps.contiguous()
|
||||
else:
|
||||
psg_out = self.model(**features, return_dict=True, output_hidden_states=True)
|
||||
p_reps = psg_out.hidden_states[-1][:, -8:, :]
|
||||
p_reps = torch.mean(p_reps, dim=1)
|
||||
if self.normlized:
|
||||
p_reps = torch.nn.functional.normalize(p_reps, dim=-1)
|
||||
return p_reps.contiguous()
|
||||
else:
|
||||
all_p_reps = []
|
||||
for sub_features in features:
|
||||
psg_out = self.model(**sub_features, return_dict=True, output_hidden_states=True)
|
||||
### modify
|
||||
p_reps = psg_out.hidden_states[-1][:, -8:, :]
|
||||
p_reps = torch.mean(p_reps, dim=1)
|
||||
all_p_reps.append(p_reps)
|
||||
all_p_reps = torch.cat(all_p_reps, 0).contiguous()
|
||||
if self.normlized:
|
||||
all_p_reps = torch.nn.functional.normalize(all_p_reps, dim=-1)
|
||||
return all_p_reps.contiguous()
|
||||
|
||||
|
||||
def compute_similarity(self, q_reps, p_reps):
|
||||
if len(p_reps.size()) == 2:
|
||||
return torch.matmul(q_reps, p_reps.transpose(0, 1))
|
||||
return torch.matmul(q_reps, p_reps.transpose(-2, -1))
|
||||
|
||||
def forward(self, query: Union[Dict[str, Tensor], List[Dict[str, Tensor]]]= None, passage: Union[Dict[str, Tensor], List[Dict[str, Tensor]]] = None):
|
||||
# torch.cuda.empty_cache()
|
||||
# if self.process_rank == 1:
|
||||
# print(query)
|
||||
q_reps = self.encode(query) # (batch_size, dim)
|
||||
p_reps = self.encode(passage) # (batch_size * num, dim)
|
||||
|
||||
if self.training:
|
||||
if self.negatives_cross_device:
|
||||
q_reps = self._dist_gather_tensor(q_reps)
|
||||
p_reps = self._dist_gather_tensor(p_reps)
|
||||
|
||||
scores = self.compute_similarity(q_reps, p_reps)
|
||||
scores = scores / self.temperature
|
||||
scores = scores.view(q_reps.size(0), -1)
|
||||
|
||||
target = torch.arange(scores.size(0), device=scores.device, dtype=torch.long)
|
||||
target = target * (p_reps.size(0) // q_reps.size(0))
|
||||
loss = self.compute_loss(scores, target) # 同批内除了正样本以外的均为负样本
|
||||
|
||||
else:
|
||||
scores = self.compute_similarity(q_reps, p_reps)
|
||||
loss = None
|
||||
|
||||
# print(loss)
|
||||
return EncoderOutput(
|
||||
loss=loss,
|
||||
scores=scores,
|
||||
q_reps=q_reps,
|
||||
p_reps=p_reps,
|
||||
)
|
||||
|
||||
def compute_loss(self, scores, target):
|
||||
return self.cross_entropy(scores, target)
|
||||
|
||||
def _dist_gather_tensor(self, t: Optional[torch.Tensor]):
|
||||
if t is None:
|
||||
return None
|
||||
t = t.contiguous()
|
||||
|
||||
all_tensors = [torch.empty_like(t) for _ in range(self.world_size)]
|
||||
dist.all_gather(all_tensors, t)
|
||||
all_tensors[self.process_rank] = t # 给当前进程的q和doc加上梯度,当前的q对其他的d，更新；当前的d对其他的q，更新
|
||||
all_tensors = torch.cat(all_tensors, dim=0)
|
||||
|
||||
return all_tensors
|
||||
|
||||
def save(self, output_dir: str):
|
||||
state_dict = self.model.state_dict()
|
||||
state_dict = type(state_dict)(
|
||||
{k: v.clone().cpu()
|
||||
for k,
|
||||
v in state_dict.items()})
|
||||
self.model.save_pretrained(output_dir, state_dict=state_dict)
|
||||
@@ -0,0 +1,125 @@
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from transformers import AutoConfig, AutoTokenizer
|
||||
from transformers import (
|
||||
HfArgumentParser,
|
||||
set_seed,
|
||||
)
|
||||
|
||||
from arguments import ModelArguments, DataArguments, \
|
||||
RetrieverTrainingArguments as TrainingArguments
|
||||
from data import TrainDatasetForEmbedding, EmbedCollator
|
||||
from modeling import BiEncoderModel
|
||||
from trainer import BiTrainer
|
||||
from load_model import get_model
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def main():
|
||||
parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
|
||||
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
|
||||
model_args: ModelArguments
|
||||
data_args: DataArguments
|
||||
training_args: TrainingArguments
|
||||
|
||||
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."
|
||||
)
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
datefmt="%m/%d/%Y %H:%M:%S",
|
||||
level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
|
||||
)
|
||||
logger.warning(
|
||||
"Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
|
||||
training_args.local_rank,
|
||||
training_args.device,
|
||||
training_args.n_gpu,
|
||||
bool(training_args.local_rank != -1),
|
||||
training_args.fp16,
|
||||
)
|
||||
logger.info("Training/evaluation parameters %s", training_args)
|
||||
logger.info("Model parameters %s", model_args)
|
||||
logger.info("Data parameters %s", data_args)
|
||||
|
||||
# Set seed
|
||||
set_seed(training_args.seed)
|
||||
|
||||
num_labels = 1
|
||||
base_model = get_model(model_args)
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
use_fast=False,
|
||||
# add_eos_token=True
|
||||
)
|
||||
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.unk_token
|
||||
tokenizer.padding_side = 'left'
|
||||
|
||||
config = AutoConfig.from_pretrained(
|
||||
model_args.config_name if model_args.config_name else model_args.model_name_or_path,
|
||||
num_labels=num_labels,
|
||||
cache_dir=model_args.cache_dir,
|
||||
token=model_args.token,
|
||||
)
|
||||
logger.info('Config: %s', config)
|
||||
|
||||
model = BiEncoderModel(model=base_model,
|
||||
tokenizer=tokenizer,
|
||||
normlized=training_args.normlized,
|
||||
negatives_cross_device=training_args.negatives_cross_device,
|
||||
temperature=training_args.temperature,
|
||||
sub_batch_size=training_args.sub_batch_size)
|
||||
# model.gradient_checkpointing_enable()
|
||||
# print(tokenizer('slalala', return_tensors='pt').to('cuda'))
|
||||
# print(base_model(**(tokenizer('slalala', return_tensors='pt'))))
|
||||
# print(base_model(**(tokenizer('slalala', return_tensors='pt').to('cuda'))))
|
||||
|
||||
if training_args.gradient_checkpointing:
|
||||
model.enable_input_require_grads()
|
||||
|
||||
train_dataset = TrainDatasetForEmbedding(args=data_args, tokenizer=tokenizer)
|
||||
|
||||
trainer = BiTrainer(
|
||||
model=model,
|
||||
args=training_args,
|
||||
train_dataset=train_dataset,
|
||||
data_collator=EmbedCollator(
|
||||
tokenizer=tokenizer,
|
||||
query_max_len=data_args.query_max_len,
|
||||
passage_max_len=data_args.passage_max_len,
|
||||
pad_to_multiple_of=8,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
sub_batch_size=training_args.sub_batch_size
|
||||
),
|
||||
tokenizer=tokenizer
|
||||
)
|
||||
|
||||
Path(training_args.output_dir).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Training
|
||||
trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
|
||||
trainer.save_model()
|
||||
# For convenience, we also re-save the tokenizer to the same directory,
|
||||
# so that you can share your model easily on huggingface.co/models =)
|
||||
if trainer.is_world_process_zero():
|
||||
tokenizer.save_pretrained(training_args.output_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,38 @@
|
||||
from transformers.trainer import *
|
||||
|
||||
|
||||
class BiTrainer(Trainer):
|
||||
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("Saving model checkpoint to %s", 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'):
|
||||
raise NotImplementedError(
|
||||
f'MODEL {self.model.__class__.__name__} '
|
||||
f'does not support save interface')
|
||||
else:
|
||||
self.model.save(output_dir)
|
||||
# if self.tokenizer is not None and self.is_world_process_zero():
|
||||
# self.tokenizer.save_pretrained(output_dir)
|
||||
|
||||
torch.save(self.args, os.path.join(output_dir, "training_args.bin"))
|
||||
|
||||
# save the checkpoint for sentence-transformers library
|
||||
# if self.is_world_process_zero():
|
||||
# save_ckpt_for_sentence_transformers(output_dir,
|
||||
# pooling_mode=self.args.sentence_pooling_method,
|
||||
# normlized=self.args.normlized)
|
||||
|
||||
def compute_loss(self, model, inputs, return_outputs=False):
|
||||
"""
|
||||
How the loss is computed by Trainer. By default, all models return the loss in the first element.
|
||||
|
||||
Subclass and override for custom behavior.
|
||||
"""
|
||||
|
||||
outputs = model(**inputs)
|
||||
loss = outputs.loss
|
||||
|
||||
return (loss, outputs) if return_outputs else loss
|
||||
@@ -0,0 +1,101 @@
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, List
|
||||
|
||||
from transformers import TrainingArguments
|
||||
|
||||
|
||||
def default_list() -> List[int]:
|
||||
return ['v_proj', 'q_proj', 'k_proj', 'gate_proj', 'down_proj', 'o_proj', 'up_proj']
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelArguments:
|
||||
"""
|
||||
Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
|
||||
"""
|
||||
|
||||
model_name_or_path: str = field(
|
||||
metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
|
||||
)
|
||||
config_name: Optional[str] = field(
|
||||
default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
|
||||
)
|
||||
tokenizer_name: Optional[str] = field(
|
||||
default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
|
||||
)
|
||||
# cache_dir: Optional[str] = field(
|
||||
# default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
|
||||
# )
|
||||
use_lora: bool = field(
|
||||
default=True,
|
||||
metadata={"help": "If passed, will use LORA (low-rank parameter-efficient training) to train the model."}
|
||||
)
|
||||
lora_rank: int = field(
|
||||
default=64,
|
||||
metadata={"help": "The rank of lora."}
|
||||
)
|
||||
lora_alpha: float = field(
|
||||
default=16,
|
||||
metadata={"help": "The alpha parameter of lora."}
|
||||
)
|
||||
lora_dropout: float = field(
|
||||
default=0.1,
|
||||
metadata={"help": "The dropout rate of lora modules."}
|
||||
)
|
||||
target_modules: List[str] = field(
|
||||
default_factory=default_list
|
||||
)
|
||||
save_merged_lora_model: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "If passed, will merge the lora modules and save the entire model."}
|
||||
)
|
||||
use_flash_attn: bool = field(
|
||||
default=True,
|
||||
metadata={"help": "If passed, will use flash attention to train the model."}
|
||||
)
|
||||
use_slow_tokenizer: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "If passed, will use a slow tokenizer (not backed by the 🤗 Tokenizers library)."}
|
||||
)
|
||||
token: str = field(
|
||||
default=""
|
||||
)
|
||||
cache_dir: str = field(
|
||||
default="./LMs"
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class DataArguments:
|
||||
cache_path: str = field(
|
||||
default='./data_dir'
|
||||
)
|
||||
|
||||
train_data: str = field(
|
||||
default='./toy_finetune_data.jsonl', metadata={"help": "Path to train data"}
|
||||
)
|
||||
|
||||
max_example_num_per_dataset: int = field(
|
||||
default=100000000, metadata={"help": "the max number of examples for each dataset"}
|
||||
)
|
||||
|
||||
cutoff_len: int = field(
|
||||
default=512,
|
||||
metadata={
|
||||
"help": "The maximum total input sequence length after tokenization for passage. Sequences longer "
|
||||
"than this will be truncated, sequences shorter will be padded."
|
||||
},
|
||||
)
|
||||
|
||||
remove_stop_words: bool = field(
|
||||
default=False
|
||||
)
|
||||
|
||||
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 PretrainTrainingArguments(TrainingArguments):
|
||||
mask: bool = field(default=True, metadata={"help": "mask the input part"})
|
||||
@@ -0,0 +1,171 @@
|
||||
import os.path
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
import re
|
||||
|
||||
import datasets
|
||||
import numpy as np
|
||||
from torch.utils.data import Dataset
|
||||
from transformers import DataCollatorWithPadding, AutoTokenizer, DataCollatorForSeq2Seq, PreTrainedTokenizer
|
||||
|
||||
from arguments import DataArguments
|
||||
from nltk.corpus import stopwords
|
||||
|
||||
class TrainDatasetForEmbedding(Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
tokenizer: PreTrainedTokenizer,
|
||||
args: DataArguments,
|
||||
):
|
||||
if os.path.isdir(args.train_data):
|
||||
train_datasets = []
|
||||
for file in os.listdir(args.train_data):
|
||||
temp_dataset = datasets.load_dataset('json', data_files=os.path.join(args.train_data, file),
|
||||
split='train')
|
||||
if len(temp_dataset) > args.max_example_num_per_dataset:
|
||||
temp_dataset = temp_dataset.select(
|
||||
random.sample(list(range(len(temp_dataset))), args.max_example_num_per_dataset))
|
||||
train_datasets.append(temp_dataset)
|
||||
self.dataset = datasets.concatenate_datasets(train_datasets)
|
||||
else:
|
||||
self.dataset = datasets.load_dataset('json', data_files=args.train_data, split='train',
|
||||
cache_dir=args.cache_path)
|
||||
|
||||
self.args = args
|
||||
self.total_len = len(self.dataset)
|
||||
|
||||
self.remove_stop_words = args.remove_stop_words
|
||||
self.stop_words = stopwords.words('english')
|
||||
self.stop_words.extend(['!', ',' ,'.' ,'?'])
|
||||
|
||||
self.max_length = args.cutoff_len
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
self.prefix = '"'
|
||||
self.suffix = ['", summarize the above passage within eight words: <s1><s2><s3><s4><s5><s6><s7><s8>',
|
||||
'", predict the following passage within eight words: <s9><s10><s11><s12><s13><s14><s15><s16>']
|
||||
self.prefix_ids = self.tokenizer(self.prefix, truncation=True, max_length=self.max_length, return_tensors=None)['input_ids']
|
||||
self.suffix_ids = self.tokenizer(self.suffix, truncation=True, max_length=self.max_length, return_tensors=None, add_special_tokens=False)['input_ids']
|
||||
|
||||
def __len__(self):
|
||||
return self.total_len
|
||||
|
||||
def __getitem__(self, item):
|
||||
prefix = self.prefix
|
||||
prefix_ids = self.prefix_ids
|
||||
|
||||
inp = self.dataset[item]['input']
|
||||
|
||||
suffix_ids_summarize = self.suffix_ids[0]
|
||||
suffix_ids_predict = self.suffix_ids[1]
|
||||
|
||||
oup_summarize = self.dataset[item]['output_summarize']
|
||||
oup_predict = self.dataset[item]['output_predict']
|
||||
|
||||
input_ids = self.tokenizer(inp,
|
||||
truncation=True,
|
||||
max_length=self.max_length - len(prefix_ids) - len(suffix_ids_summarize) - len(suffix_ids_predict),
|
||||
padding=False,
|
||||
return_tensors=None,
|
||||
add_special_tokens=False)
|
||||
result = dict()
|
||||
result['input_ids'] = prefix_ids + input_ids['input_ids'] + suffix_ids_summarize + suffix_ids_predict
|
||||
result['attention_mask'] = [1] * len(result['input_ids'])
|
||||
result['labels'] = [-100] * len(prefix_ids) + result['input_ids'][
|
||||
len(prefix_ids) : len(result['input_ids']) - len(suffix_ids_summarize) - len(suffix_ids_predict)] + [
|
||||
-100] * (len(suffix_ids_summarize) + len(suffix_ids_predict))
|
||||
if self.remove_stop_words:
|
||||
# oup = re.sub(r'[\d\W_]+', ' ', oup)
|
||||
oup_summarize = re.sub(r'[\W_]+', ' ', oup_summarize)
|
||||
oup_summarize = ' '.join([word for word in oup_summarize.split() if word.lower() not in self.stop_words])
|
||||
|
||||
oup_predict = re.sub(r'[\W_]+', ' ', oup_predict)
|
||||
oup_predict = ' '.join([word for word in oup_predict.split() if word.lower() not in self.stop_words])
|
||||
return result, oup_summarize, oup_predict
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmbedCollator(DataCollatorForSeq2Seq):
|
||||
"""
|
||||
Wrapper that does conversion from List[Tuple[encode_qry, encode_psg]] to List[qry], List[psg]
|
||||
and pass batch separately to the actual collator.
|
||||
Abstract out data detail for the model.
|
||||
"""
|
||||
cutoff_len: int = 512
|
||||
|
||||
def __call__(self, features, return_tensors='pt'):
|
||||
if return_tensors is None:
|
||||
return_tensors = self.return_tensors
|
||||
|
||||
inputs = []
|
||||
outputs_summarize = []
|
||||
outputs_predict = []
|
||||
for e in features:
|
||||
inputs.append(e[0])
|
||||
outputs_summarize.append(e[1])
|
||||
outputs_predict.append(e[2])
|
||||
|
||||
labels = [feature["labels"] for feature in inputs] if "labels" in inputs[0].keys() else None
|
||||
# We have to pad the labels before calling `tokenizer.pad` as this method won't pad them and needs them of the
|
||||
# same length to return tensors.
|
||||
if labels is not None:
|
||||
max_label_length = max(len(l) for l in labels)
|
||||
if self.pad_to_multiple_of is not None:
|
||||
max_label_length = (
|
||||
(max_label_length + self.pad_to_multiple_of - 1)
|
||||
// self.pad_to_multiple_of
|
||||
* self.pad_to_multiple_of
|
||||
)
|
||||
|
||||
padding_side = self.tokenizer.padding_side
|
||||
for feature in inputs:
|
||||
remainder = [self.label_pad_token_id] * (max_label_length - len(feature["labels"]))
|
||||
if isinstance(feature["labels"], list):
|
||||
feature["labels"] = (
|
||||
feature["labels"] + remainder if padding_side == "right" else remainder + feature["labels"]
|
||||
)
|
||||
elif padding_side == "right":
|
||||
feature["labels"] = np.concatenate([feature["labels"], remainder]).astype(np.int64)
|
||||
else:
|
||||
feature["labels"] = np.concatenate([remainder, feature["labels"]]).astype(np.int64)
|
||||
|
||||
inputs = self.tokenizer.pad(
|
||||
inputs,
|
||||
padding=self.padding,
|
||||
max_length=self.cutoff_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=return_tensors,
|
||||
)
|
||||
|
||||
# prepare decoder_input_ids
|
||||
if (
|
||||
labels is not None
|
||||
and self.model is not None
|
||||
and hasattr(self.model, "prepare_decoder_input_ids_from_labels")
|
||||
):
|
||||
decoder_input_ids = self.model.prepare_decoder_input_ids_from_labels(labels=inputs["labels"])
|
||||
inputs["decoder_input_ids"] = decoder_input_ids
|
||||
|
||||
outputs_summarize_collated = self.tokenizer(
|
||||
outputs_summarize,
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=self.cutoff_len,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
outputs_predict_collated = self.tokenizer(
|
||||
outputs_predict,
|
||||
padding=True,
|
||||
truncation=True,
|
||||
max_length=self.cutoff_len,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
return {"input_ids": inputs['input_ids'],
|
||||
"attention_mask": inputs['attention_mask'],
|
||||
"labels": inputs['labels'],
|
||||
"output_summarize_ids": outputs_summarize_collated['input_ids'],
|
||||
"output_predict_ids": outputs_predict_collated['input_ids']}
|
||||
@@ -0,0 +1,50 @@
|
||||
|
||||
from transformers import AutoConfig, AutoModelForCausalLM
|
||||
from peft import LoraConfig, TaskType, get_peft_model
|
||||
from modeling import PreLlamaModel
|
||||
|
||||
def get_model(model_args, use_gradient_checkpointing: bool = False):
|
||||
if model_args.config_name:
|
||||
config = AutoConfig.from_pretrained(model_args.config_name,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
)
|
||||
elif model_args.model_name_or_path:
|
||||
config = AutoConfig.from_pretrained(model_args.model_name_or_path,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
"You are instantiating a new config instance from scratch. This is not supported by this script."
|
||||
)
|
||||
if use_gradient_checkpointing:
|
||||
config.use_cache = False
|
||||
|
||||
if model_args.model_name_or_path:
|
||||
model = PreLlamaModel.from_pretrained(
|
||||
model_args.model_name_or_path,
|
||||
use_flash_attention_2=True if model_args.use_flash_attn else False,
|
||||
attn_implementation='sdpa',
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
from_tf=bool(".ckpt" in model_args.model_name_or_path),
|
||||
config=config,
|
||||
)
|
||||
else:
|
||||
print("Training new model from scratch")
|
||||
model = model_args.from_config(config)
|
||||
|
||||
if model_args.use_lora:
|
||||
peft_config = LoraConfig(
|
||||
task_type="CAUSAL_LM",
|
||||
inference_mode=False,
|
||||
r=model_args.lora_rank,
|
||||
target_modules=model_args.target_modules,
|
||||
lora_alpha=model_args.lora_alpha,
|
||||
lora_dropout=model_args.lora_dropout
|
||||
)
|
||||
model = get_peft_model(model, peft_config)
|
||||
model.print_trainable_parameters()
|
||||
|
||||
return model
|
||||
@@ -0,0 +1,441 @@
|
||||
import sys
|
||||
from typing import Optional, List, Union, Tuple
|
||||
|
||||
import torch
|
||||
import copy
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
from torch.nn import CrossEntropyLoss
|
||||
from transformers import LlamaForCausalLM, LlamaPreTrainedModel, LlamaConfig, AutoModel
|
||||
from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutputWithPast
|
||||
from transformers.models.idefics.modeling_idefics import LLAMA_INPUTS_DOCSTRING, _CONFIG_FOR_DOC
|
||||
from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaRMSNorm, LlamaModel
|
||||
from transformers.utils import add_start_docstrings_to_model_forward, replace_return_docstrings, logging
|
||||
from transformers.cache_utils import Cache, DynamicCache, StaticCache
|
||||
from transformers.modeling_attn_mask_utils import AttentionMaskConverter
|
||||
import torch.distributed as dist
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
class NewLlamaModel(LlamaModel):
|
||||
add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
) -> Union[Tuple, BaseModelOutputWithPast]:
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
use_cache = use_cache if use_cache is not None else self.config.use_cache
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
if self.config._attn_implementation == "flash_attention_2":
|
||||
raise ValueError(
|
||||
"You can not use flash attention to pretrain"
|
||||
)
|
||||
|
||||
if (input_ids is None) ^ (inputs_embeds is not None):
|
||||
raise ValueError(
|
||||
"You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
|
||||
)
|
||||
|
||||
if self.gradient_checkpointing and self.training and use_cache:
|
||||
logger.warning_once(
|
||||
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
|
||||
)
|
||||
use_cache = False
|
||||
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds = self.embed_tokens(input_ids)
|
||||
|
||||
return_legacy_cache = False
|
||||
if use_cache and not isinstance(past_key_values, Cache): # kept for BC (non `Cache` `past_key_values` inputs)
|
||||
return_legacy_cache = True
|
||||
past_key_values = DynamicCache.from_legacy_cache(past_key_values)
|
||||
|
||||
if cache_position is None:
|
||||
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
|
||||
cache_position = torch.arange(
|
||||
past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
|
||||
)
|
||||
|
||||
summarize_suffix_ids = [9162, 19138, 675, 278, 2038, 13382, 2629, 9475, 3838, 29901, 29871,
|
||||
32008, 32011, 32004, 32013, 32007, 32005, 32002, 32014]
|
||||
predict_suffix_ids = [9162, 8500, 278, 1494, 13382, 2629, 9475, 3838, 29901, 29871, 32000,
|
||||
32009, 32012, 32001, 32010, 32003, 32006, 32015]
|
||||
|
||||
if position_ids is None:
|
||||
position_ids = cache_position.unsqueeze(0)
|
||||
for i in range(len(position_ids)):
|
||||
position_ids[i][-len(predict_suffix_ids):] = copy.deepcopy(position_ids[i][
|
||||
-len(summarize_suffix_ids) - len(predict_suffix_ids): -len(
|
||||
summarize_suffix_ids)])
|
||||
|
||||
causal_mask = self._update_causal_mask(
|
||||
attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
|
||||
)
|
||||
|
||||
causal_mask[:,
|
||||
:,
|
||||
-len(predict_suffix_ids) :,
|
||||
-len(predict_suffix_ids) - len(summarize_suffix_ids): -len(predict_suffix_ids),
|
||||
] = torch.finfo(inputs_embeds.dtype).min
|
||||
|
||||
# embed positions
|
||||
hidden_states = inputs_embeds
|
||||
|
||||
# decoder layers
|
||||
all_hidden_states = () if output_hidden_states else None
|
||||
all_self_attns = () if output_attentions else None
|
||||
next_decoder_cache = None
|
||||
|
||||
for decoder_layer in self.layers:
|
||||
if output_hidden_states:
|
||||
all_hidden_states += (hidden_states,)
|
||||
|
||||
if self.gradient_checkpointing and self.training:
|
||||
layer_outputs = self._gradient_checkpointing_func(
|
||||
decoder_layer.__call__,
|
||||
hidden_states,
|
||||
causal_mask,
|
||||
position_ids,
|
||||
past_key_values,
|
||||
output_attentions,
|
||||
use_cache,
|
||||
cache_position,
|
||||
)
|
||||
else:
|
||||
layer_outputs = decoder_layer(
|
||||
hidden_states,
|
||||
attention_mask=causal_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_value=past_key_values,
|
||||
output_attentions=output_attentions,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
)
|
||||
|
||||
hidden_states = layer_outputs[0]
|
||||
|
||||
if use_cache:
|
||||
next_decoder_cache = layer_outputs[2 if output_attentions else 1]
|
||||
|
||||
if output_attentions:
|
||||
all_self_attns += (layer_outputs[1],)
|
||||
|
||||
hidden_states = self.norm(hidden_states)
|
||||
|
||||
# add hidden states from the last decoder layer
|
||||
if output_hidden_states:
|
||||
all_hidden_states += (hidden_states,)
|
||||
|
||||
next_cache = next_decoder_cache if use_cache else None
|
||||
if return_legacy_cache:
|
||||
next_cache = next_cache.to_legacy_cache()
|
||||
|
||||
if not return_dict:
|
||||
return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state=hidden_states,
|
||||
past_key_values=next_cache,
|
||||
hidden_states=all_hidden_states,
|
||||
attentions=all_self_attns,
|
||||
)
|
||||
|
||||
def _update_causal_mask(
|
||||
self,
|
||||
attention_mask: torch.Tensor,
|
||||
input_tensor: torch.Tensor,
|
||||
cache_position: torch.Tensor,
|
||||
past_key_values: Cache,
|
||||
output_attentions: bool,
|
||||
):
|
||||
# TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static
|
||||
# KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes.
|
||||
# (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using
|
||||
# `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/29114
|
||||
|
||||
if self.config._attn_implementation == "flash_attention_2":
|
||||
if attention_mask is not None and 0.0 in attention_mask:
|
||||
return attention_mask
|
||||
return None
|
||||
|
||||
# For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
|
||||
# order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
|
||||
# to infer the attention mask.
|
||||
past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
|
||||
using_static_cache = isinstance(past_key_values, StaticCache)
|
||||
|
||||
# When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
|
||||
# if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
|
||||
# if AttentionMaskConverter._ignore_causal_mask_sdpa(
|
||||
# attention_mask,
|
||||
# inputs_embeds=input_tensor,
|
||||
# past_key_values_length=past_seen_tokens,
|
||||
# is_training=self.training,
|
||||
# ):
|
||||
# return None
|
||||
|
||||
dtype, device = input_tensor.dtype, input_tensor.device
|
||||
min_dtype = torch.finfo(dtype).min
|
||||
sequence_length = input_tensor.shape[1]
|
||||
if using_static_cache:
|
||||
target_length = past_key_values.get_max_length()
|
||||
else:
|
||||
target_length = (
|
||||
attention_mask.shape[-1]
|
||||
if isinstance(attention_mask, torch.Tensor)
|
||||
else past_seen_tokens + sequence_length + 1
|
||||
)
|
||||
|
||||
if attention_mask is not None and attention_mask.dim() == 4:
|
||||
# in this case we assume that the mask comes already in inverted form and requires no inversion or slicing
|
||||
if attention_mask.max() != 0:
|
||||
raise ValueError("Custom 4D attention mask should be passed in inverted form with max==0`")
|
||||
causal_mask = attention_mask
|
||||
else:
|
||||
causal_mask = torch.full(
|
||||
(sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
|
||||
)
|
||||
if sequence_length != 1:
|
||||
causal_mask = torch.triu(causal_mask, diagonal=1)
|
||||
causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
|
||||
causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
|
||||
if attention_mask is not None:
|
||||
causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
|
||||
mask_length = attention_mask.shape[-1]
|
||||
padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
|
||||
padding_mask = padding_mask == 0
|
||||
causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
|
||||
padding_mask, min_dtype
|
||||
)
|
||||
if (
|
||||
self.config._attn_implementation == "sdpa"
|
||||
and attention_mask is not None
|
||||
and attention_mask.device.type == "cuda"
|
||||
and not output_attentions
|
||||
):
|
||||
# Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
|
||||
# using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
|
||||
# Details: https://github.com/pytorch/pytorch/issues/110213
|
||||
causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)
|
||||
|
||||
return causal_mask
|
||||
|
||||
class PreLlamaModel(LlamaForCausalLM):
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
self.model = NewLlamaModel(config)
|
||||
self.log_softmax = torch.nn.LogSoftmax(dim=1)
|
||||
|
||||
"""
|
||||
prompt type1: "{}", summarize the above passage within eight words: <s1><s2><s3><s4><s5><s6><s7><s8>
|
||||
token ids: [376, ..., 9162, 19138, 675, 278, 2038, 13382, 2629, 9475, 3838, 29901, 29871,
|
||||
32008, 32011, 32004, 32013, 32007, 32005, 32002, 32014]
|
||||
|
||||
prompt type2: "{}", predict the following passage within eight words: <s9><s10><s11><s12><s13><s14><s15><s16>
|
||||
token ids: [376, ..., 8500, 278, 1494, 13382, 2629, 9475, 3838, 29901, 29871, 32000, 32009,
|
||||
32012, 32001, 32010, 32003, 32006, 32015]
|
||||
|
||||
[9162, 19138, 675, 278, 2038, 13382, 2629, 9475, 3838, 29901,
|
||||
29871, 32008, 32011, 32004, 32013, 32007, 32005, 32002, 32014,
|
||||
8500, 278, 1494, 13382, 2629, 9475, 3838, 29901, 29871, 32000,
|
||||
32009, 32012, 32001, 32010, 32003, 32006, 32015]
|
||||
|
||||
Maybe only one of them will appear, or both may appear. We consider all possibilities here.
|
||||
"""
|
||||
|
||||
self.summarize_prompt_ids = [9162, 19138, 675, 278, 2038, 13382, 2629, 9475, 3838, 29901, 29871,
|
||||
32008, 32011, 32004, 32013, 32007, 32005, 32002, 32014]
|
||||
self.predict_prompt_ids = [9162, 8500, 278, 1494, 13382, 2629, 9475, 3838, 29901, 29871, 32000,
|
||||
32009, 32012, 32001, 32010, 32003, 32006, 32015]
|
||||
|
||||
@add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
|
||||
@replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
|
||||
def forward(
|
||||
self,
|
||||
input_ids: torch.LongTensor = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
output_summarize_ids: Optional[torch.LongTensor] = None,
|
||||
output_predict_ids: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
) -> Union[Tuple, CausalLMOutputWithPast]:
|
||||
r"""
|
||||
Args:
|
||||
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
|
||||
Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
|
||||
config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
|
||||
(masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
|
||||
|
||||
Returns:
|
||||
|
||||
Example:
|
||||
|
||||
```python
|
||||
>>> from transformers import AutoTokenizer, LlamaForCausalLM
|
||||
|
||||
>>> model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
|
||||
>>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
|
||||
|
||||
>>> prompt = "Hey, are you conscious? Can you talk to me?"
|
||||
>>> inputs = tokenizer(prompt, return_tensors="pt")
|
||||
|
||||
>>> # Generate
|
||||
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
|
||||
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
|
||||
"Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
|
||||
```"""
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
|
||||
outputs = self.model(
|
||||
input_ids=input_ids,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
cache_position=cache_position,
|
||||
)
|
||||
|
||||
hidden_states = outputs[0]
|
||||
if self.config.pretraining_tp > 1:
|
||||
lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
|
||||
logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]
|
||||
logits = torch.cat(logits, dim=-1)
|
||||
else:
|
||||
logits = self.lm_head(hidden_states)
|
||||
logits = logits.float()
|
||||
|
||||
ar_loss = None
|
||||
if labels is not None:
|
||||
# Shift so that tokens < n predict n
|
||||
shift_logits = logits[..., :-1, :].contiguous()
|
||||
shift_labels = labels[..., 1:].contiguous()
|
||||
# Flatten the tokens
|
||||
loss_fct = CrossEntropyLoss()
|
||||
shift_logits = shift_logits.view(-1, self.config.vocab_size)
|
||||
shift_labels = shift_labels.view(-1)
|
||||
# Enable model parallelism
|
||||
shift_labels = shift_labels.to(shift_logits.device)
|
||||
ar_loss = loss_fct(shift_logits, shift_labels)
|
||||
|
||||
"""
|
||||
prompt: ", summarize the above passage within eight words: <s1><s2><s3><s4><s5><s6><s7><s8>", predict the following passage within eight words: <s9><s10><s11><s12><s13><s14><s15><s16>
|
||||
token ids: [9162, 19138, 675, 278, 2038, 13382, 2629, 9475, 3838, 29901,
|
||||
29871, 32008, 32011, 32004, 32013, 32007, 32005, 32002, 32014,
|
||||
9162, 8500, 278, 1494, 13382, 2629, 9475, 3838, 29901, 29871,
|
||||
32000, 32009, 32012, 32001, 32010, 32003, 32006, 32015]
|
||||
token ids[-26:-18] —— <s1><s2><s3><s4><s5><s6><s7><s8>
|
||||
token ids[-8:] —— <s9><s10><s11><s12><s13><s14><s15><s16>
|
||||
"""
|
||||
bow_summarize_loss = 0
|
||||
if output_summarize_ids is not None:
|
||||
special_logits = logits[:, -len(self.predict_prompt_ids) - 8:-len(self.predict_prompt_ids), :]
|
||||
special_logits, _ = torch.max(special_logits, dim=1)
|
||||
bow_summarize_loss = 0
|
||||
possibility = self.log_softmax(special_logits)
|
||||
batch_num = 0
|
||||
for p, temp_output_ids in zip(possibility, output_summarize_ids):
|
||||
unique_useful_ids = torch.unique(temp_output_ids[temp_output_ids > 2])
|
||||
if len(unique_useful_ids) > 0:
|
||||
bow_summarize_loss -= torch.mean(p[unique_useful_ids])
|
||||
batch_num += 1
|
||||
if batch_num > 0:
|
||||
bow_summarize_loss /= batch_num
|
||||
bow_summarize_loss /= 10
|
||||
|
||||
bow_predict_loss = 0
|
||||
if output_predict_ids is not None:
|
||||
special_logits = logits[:, -8:, :]
|
||||
special_logits, _ = torch.max(special_logits, dim=1)
|
||||
bow_predict_loss = 0
|
||||
possibility = self.log_softmax(special_logits)
|
||||
batch_num = 0
|
||||
for p, temp_output_ids in zip(possibility, output_predict_ids):
|
||||
unique_useful_ids = torch.unique(temp_output_ids[temp_output_ids > 2])
|
||||
if len(unique_useful_ids) > 0:
|
||||
bow_predict_loss -= torch.mean(p[unique_useful_ids])
|
||||
batch_num += 1
|
||||
if batch_num > 0:
|
||||
bow_predict_loss /= batch_num
|
||||
bow_predict_loss /= 10
|
||||
|
||||
if bow_summarize_loss > 0 and bow_predict_loss > 0:
|
||||
bow_loss = (bow_summarize_loss + bow_predict_loss) / 2
|
||||
elif bow_summarize_loss > 0:
|
||||
bow_loss = bow_summarize_loss
|
||||
elif bow_predict_loss > 0:
|
||||
bow_loss = bow_predict_loss
|
||||
else:
|
||||
bow_loss = None
|
||||
|
||||
if ar_loss is not None and bow_loss is not None:
|
||||
loss = ar_loss + bow_loss
|
||||
elif ar_loss is None:
|
||||
loss = bow_loss
|
||||
else:
|
||||
loss = ar_loss
|
||||
|
||||
if not return_dict:
|
||||
output = (logits,) + outputs[1:]
|
||||
return (loss,) + output if loss is not None else output
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
|
||||
class PreModel(nn.Module):
|
||||
def __init__(self,
|
||||
model: AutoModel = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.model = model
|
||||
|
||||
def gradient_checkpointing_enable(self, **kwargs):
|
||||
self.model.gradient_checkpointing_enable(**kwargs)
|
||||
|
||||
def enable_input_require_grads(self, **kwargs):
|
||||
self.model.enable_input_require_grads(**kwargs)
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
return self.model(*args, **kwargs)
|
||||
|
||||
def save(self, output_dir: str):
|
||||
state_dict = self.model.state_dict()
|
||||
state_dict = type(state_dict)(
|
||||
{k: v.clone().cpu()
|
||||
for k,
|
||||
v in state_dict.items()})
|
||||
self.model.save_pretrained(output_dir, state_dict=state_dict)
|
||||
@@ -0,0 +1,132 @@
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import transformers
|
||||
|
||||
from transformers import AutoTokenizer, HfArgumentParser, set_seed, AutoConfig, Trainer
|
||||
|
||||
from arguments import ModelArguments, DataArguments, \
|
||||
PretrainTrainingArguments as TrainingArguments
|
||||
from data import TrainDatasetForEmbedding, EmbedCollator
|
||||
from load_model import get_model
|
||||
from modeling import PreModel
|
||||
from trainer import PreTrainer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def main():
|
||||
parser = HfArgumentParser((ModelArguments, DataArguments, TrainingArguments))
|
||||
model_args, data_args, training_args = parser.parse_args_into_dataclasses()
|
||||
model_args: ModelArguments
|
||||
data_args: DataArguments
|
||||
training_args: TrainingArguments
|
||||
|
||||
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."
|
||||
)
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
datefmt="%m/%d/%Y %H:%M:%S",
|
||||
level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
|
||||
)
|
||||
logger.warning(
|
||||
"Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
|
||||
training_args.local_rank,
|
||||
training_args.device,
|
||||
training_args.n_gpu,
|
||||
bool(training_args.local_rank != -1),
|
||||
training_args.fp16,
|
||||
)
|
||||
logger.info("Training/evaluation parameters %s", training_args)
|
||||
logger.info("Model parameters %s", model_args)
|
||||
logger.info("Data parameters %s", data_args)
|
||||
|
||||
# Set seed
|
||||
set_seed(training_args.seed)
|
||||
|
||||
num_labels = 1
|
||||
|
||||
config = AutoConfig.from_pretrained(
|
||||
model_args.config_name if model_args.config_name else model_args.model_name_or_path,
|
||||
num_labels=num_labels,
|
||||
cache_dir=model_args.cache_dir,
|
||||
)
|
||||
logger.info('Config: %s', config)
|
||||
|
||||
model = get_model(model_args, training_args.gradient_checkpointing)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
# model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
|
||||
model_args.model_name_or_path,
|
||||
token=model_args.token,
|
||||
cache_dir=model_args.cache_dir,
|
||||
use_fast=False
|
||||
)
|
||||
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.unk_token
|
||||
# special_tokens_dict = {'additional_special_tokens': ['<s1>', '<s2>', '<s3>', '<s4>',
|
||||
# '<s5>', '<s6>', '<s7>', '<s8>',
|
||||
# '<s9>', '<s10>', '<s11>', '<s12>',
|
||||
# '<s13>', '<s14>', '<s15>', '<s16>', ]}
|
||||
# tokenizer.add_special_tokens(special_tokens_dict)
|
||||
tokenizer.padding_side = "left" # Allow batched inference
|
||||
print(tokenizer)
|
||||
|
||||
special_tokens = ['<s1>', '<s2>', '<s3>', '<s4>',
|
||||
'<s5>', '<s6>', '<s7>', '<s8>',
|
||||
'<s9>', '<s10>', '<s11>', '<s12>',
|
||||
'<s13>', '<s14>', '<s15>', '<s16>', ]
|
||||
current_vocab = tokenizer.get_vocab()
|
||||
tokens_to_add = [token for token in special_tokens if token not in current_vocab]
|
||||
if tokens_to_add:
|
||||
special_tokens_dict = {'additional_special_tokens': tokens_to_add}
|
||||
tokenizer.add_special_tokens(special_tokens_dict)
|
||||
model.resize_token_embeddings(len(tokenizer))
|
||||
print(tokenizer)
|
||||
|
||||
if training_args.gradient_checkpointing:
|
||||
model.enable_input_require_grads()
|
||||
|
||||
train_dataset = TrainDatasetForEmbedding(tokenizer, args=data_args)
|
||||
|
||||
trainer = Trainer(
|
||||
model=model,
|
||||
train_dataset=train_dataset,
|
||||
args=training_args,
|
||||
data_collator=EmbedCollator(
|
||||
tokenizer=tokenizer,
|
||||
cutoff_len=data_args.cutoff_len,
|
||||
pad_to_multiple_of=8,
|
||||
return_tensors="pt",
|
||||
padding=True
|
||||
)
|
||||
)
|
||||
|
||||
Path(training_args.output_dir).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if trainer.is_world_process_zero():
|
||||
tokenizer.save_pretrained(training_args.output_dir)
|
||||
|
||||
trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
|
||||
trainer.save_model()
|
||||
if training_args.deepspeed:
|
||||
trainer.deepspeed.save_checkpoint(training_args.output_dir)
|
||||
|
||||
print("\n If there's a warning about missing keys above, please disregard :)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,31 @@
|
||||
from transformers.trainer import *
|
||||
|
||||
class PreTrainer(Trainer):
|
||||
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("Saving model checkpoint to %s", 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'):
|
||||
raise NotImplementedError(
|
||||
f'MODEL {self.model.__class__.__name__} '
|
||||
f'does not support save interface')
|
||||
else:
|
||||
self.model.save(output_dir)
|
||||
# if self.tokenizer is not None and self.is_world_process_zero():
|
||||
# self.tokenizer.save_pretrained(output_dir)
|
||||
|
||||
torch.save(self.args, os.path.join(output_dir, "training_args.bin"))
|
||||
|
||||
def compute_loss(self, model, inputs, return_outputs=False):
|
||||
"""
|
||||
How the loss is computed by Trainer. By default, all models return the loss in the first element.
|
||||
|
||||
Subclass and override for custom behavior.
|
||||
"""
|
||||
|
||||
outputs = model(**inputs)
|
||||
loss = outputs.loss
|
||||
|
||||
return (loss, outputs) if return_outputs else loss
|
||||
@@ -0,0 +1,50 @@
|
||||
{
|
||||
"zero_optimization": {
|
||||
"stage": 1,
|
||||
"reduce_bucket_size": 5e8
|
||||
},
|
||||
|
||||
"fp16": {
|
||||
"enabled": "auto",
|
||||
"loss_scale": 0,
|
||||
"initial_scale_power": 10,
|
||||
"loss_scale_window": 1000,
|
||||
"hysteresis": 2,
|
||||
"min_loss_scale": 1
|
||||
},
|
||||
"bf16": {
|
||||
"enabled": "auto",
|
||||
"loss_scale": 0,
|
||||
"initial_scale_power": 10,
|
||||
"loss_scale_window": 1000,
|
||||
"hysteresis": 2,
|
||||
"min_loss_scale": 1
|
||||
},
|
||||
"optimizer": {
|
||||
"type": "AdamW",
|
||||
"params": {
|
||||
"lr": "auto",
|
||||
"betas": "auto",
|
||||
"eps": "auto",
|
||||
"weight_decay": "auto",
|
||||
"torch_adam": true
|
||||
}
|
||||
},
|
||||
|
||||
"scheduler": {
|
||||
"type": "WarmupDecayLR",
|
||||
"params": {
|
||||
"warmup_min_lr": "auto",
|
||||
"warmup_max_lr": "auto",
|
||||
"warmup_num_steps": "auto",
|
||||
"total_num_steps": "auto"
|
||||
}
|
||||
},
|
||||
|
||||
"gradient_accumulation_steps": "auto",
|
||||
"gradient_clipping": "auto",
|
||||
"steps_per_print": 1000,
|
||||
"train_batch_size": "auto",
|
||||
"train_micro_batch_size_per_gpu": "auto",
|
||||
"wall_clock_breakdown": false
|
||||
}
|
||||
Reference in New Issue
Block a user