text embedding
# coding=utf8

import os
import logging
import yaml
import torch
from torch.utils.data import  DataLoader
from os.path import join
from copy import deepcopy
from transformers import AutoTokenizer, AutoModel, BertModel
from transformers import TrainingArguments, Trainer
import shutil

os.environ["WANDB_DISABLED"] = "true"
transformers_logger = logging.getLogger("transformers")
transformers_logger.setLevel(logging.WARNING)
import torch.nn.functional as F
from loguru import logger

from src import (
    InBatchDataSet,
    in_batch_collate_fn,
    PairDataSet,
    pair_collate_fn,
    VecDataSet,
    get_mean_params, SaveModelCallBack, cosent_loss
)


class MyTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        def get_vecs_e5(ipt):
            attention_mask = ipt["attention_mask"]
            model_output = model(**ipt)
            last_hidden = model_output.last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
            vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
            vectors = F.normalize(vectors, 2.0, dim=1)
            return vectors

        def get_vecs_bge(ipt):
            # print("input_ids.shape", ipt["input_ids"].shape)
            token_embeddings = self.model(**ipt)[0]
            vectors = token_embeddings[:, 0, :].squeeze(1)  # bsz*h
            vectors = F.normalize(vectors, 2.0, dim=1)
            return vectors

        # print(f"len(inputs){inputs[0]}", len(inputs))
        # Step1 计算inbatch loss
        q_num = inputs[-1]
        name = inputs[0]
        inputs = inputs[1:-1]
        in_batch_loss, pair_loss = torch.tensor(0.0), torch.tensor(0.0)
        if "in_batch" in name:
            if model_name in ["e5", "piccolo"]:
                vectors = [get_vecs_e5(ipt) for ipt in inputs]
            elif model_name in ["bge", "simbert", "simbert_hp"]:
                vectors = [get_vecs_bge(ipt) for ipt in inputs]
            else:
                raise NotImplementedError()
            vectors = torch.cat(vectors, dim=0)
            vecs1, vecs2 = vectors[:q_num, :], vectors[q_num:, :]
            logits = torch.mm(vecs1, vecs2.t())
            print("logits.shape", logits.shape)
            LABEL = torch.LongTensor(list(range(q_num))).to(vectors.device)
            in_batch_loss = F.cross_entropy(logits * in_batch_ratio, LABEL)

        # Step2 计算pair loss
        elif "pair" in name:
            neg_pos_idxs = inputs[-1]
            inputs = inputs[:-1]
            if model_name in ["e5", "piccolo"]:
                vectors = [get_vecs_e5(ipt) for ipt in inputs]
            elif model_name in ["bge", "simbert", "simbert_hp"]:
                vectors = [get_vecs_bge(ipt) for ipt in inputs]
            else:
                raise NotImplementedError()
            vectors = torch.cat(vectors, dim=0)
            vecs1, vecs2 = vectors[:q_num, :], vectors[q_num:, :]

            pred_sims = F.cosine_similarity(vecs1, vecs2)
            # print(name, pred_sims.shape)
            pair_loss = cosent_loss(
                neg_pos_idxs=neg_pos_idxs,
                pred_sims=pred_sims,
                cosent_ratio=cosent_ratio,
                zero_data=torch.tensor([0.0]).to(vectors.device)
            )
        # Step3 计算 ewc loss
        losses = []
        for n, p in model.named_parameters():
            # 每个参数都有mean和fisher
            mean = original_weight[n.replace("module.", "")]
            if "position_embeddings.weight" in n:
                print(p.shape, mean.shape)
                losses.append(
                    ((p - mean)[:512, :] ** 2).sum()
                )
            else:
                losses.append(
                    ((p - mean) ** 2).sum()
                )
        ewc_loss = sum(losses)

        final_loss = in_batch_loss + pair_loss
        if "ewc" in train_method:
            final_loss += (ewc_loss * ewc_ratio)
        if "in_batch" in name:
            logger.info(
                f"step-{self.state.global_step}, {name}-loss:{in_batch_loss.item()}, ewc_loss:{ewc_loss.item()}"
            )
        else:
            logger.info(
                f"step-{self.state.global_step}, {name}-loss:{pair_loss.item()}, ewc_loss:{ewc_loss.item()}"
            )
        return (final_loss, None) if return_outputs else final_loss


if __name__ == "__main__":
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    pair_label_map = {
        "0": 0,
        "1": 1,
        "contradiction": 2,
        "neutral": 3,
        "entailment": 4,
    }

    MODEL_NAME_INFO = {
        "e5": [AutoModel, AutoTokenizer, in_batch_collate_fn],
        "bge": [AutoModel, AutoTokenizer, in_batch_collate_fn],
        "piccolo": [BertModel, AutoTokenizer, in_batch_collate_fn],
        "simbert_hp": [BertModel, AutoTokenizer, in_batch_collate_fn],
    }

    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    # 读取参数并赋值
    config_path = 'conf.yml'
    with open(config_path, "r", encoding="utf8") as fr:
        conf = yaml.safe_load(fr)

    # args of hf trainer
    hf_args = deepcopy(conf["train_args"])
    in_batch_bsz = conf["in_batch_bsz"]
    pair_bsz = conf["pair_bsz"]
    if hf_args.get("deepspeed") and conf["use_deepspeed"]:
        hf_args["deepspeed"]["gradient_accumulation_steps"] = hf_args["gradient_accumulation_steps"]
        hf_args["deepspeed"]["train_micro_batch_size_per_gpu"] = hf_args["per_device_train_batch_size"]
        hf_args["deepspeed"]["optimizer"]["params"]["lr"] = hf_args["learning_rate"]
    else:
        hf_args.pop("deepspeed", None)

    model_name = conf["model_name"]

    grad_checkpoint = hf_args["gradient_checkpointing"]
    task_name = conf["task_name"]
    in_batch_train_paths = conf["in_batch_train_paths"]
    pair_train_paths = conf["pair_train_paths"]
    loader_idxs = conf["loader_idxs"]
    max_length = conf["max_length"]
    model_dir = conf["model_dir"]
    train_method = conf["train_method"]
    ewc_ratio = conf["ewc_ratio"]
    cosent_ratio = conf["cosent_ratio"]
    in_batch_ratio = conf["in_batch_ratio"]
    hard_neg_ratio = conf["hard_neg_ratio"]

    output_dir = hf_args["output_dir"]

    # 构建训练输出目录
    if world_size == 1 and conf["auto_ouput_dir"]:
        version = 1
        save_dir = join(output_dir,
                        f"{model_name}_{task_name}_bsz{in_batch_bsz}_len{max_length}_{train_method}_v{version}")
        while os.path.exists(save_dir):
            version += 1
            save_dir = join(output_dir,
                            f"{model_name}_{task_name}_bsz{in_batch_bsz}_len{max_length}_{train_method}_v{version}")
        output_dir = save_dir
        hf_args["output_dir"] = output_dir
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
    else:
        if not os.path.exists(hf_args["output_dir"]):
            os.makedirs(hf_args["output_dir"], exist_ok=True)
    # 拷贝 config
    if local_rank == 0:
        shutil.copy('conf.yml', os.path.join(output_dir, "train_config.yml"))
        # 初始化log
        logger.add(join(output_dir, "train_log.txt"), level="INFO", compression="zip", rotation="500 MB",
                   format="{message}")
    # in-batch 数据集
    in_batch_data_loaders = []
    if in_batch_train_paths:
        for data_name, data_paths in in_batch_train_paths.items():
            logger.info(f"添加数据迭代器，data_name:{data_name}, data_paths:{data_paths}")
            in_batch_data_loaders.append(
                DataLoader(
                    dataset=InBatchDataSet(data_paths=data_paths, data_name=data_name, model_name=model_name),
                    shuffle=True,
                    collate_fn=lambda x: in_batch_collate_fn(x, tokenizer, max_length),
                    drop_last=True,
                    batch_size=in_batch_bsz,
                    num_workers=2
                )
            )

    # pair对数据集
    pair_data_loaders = []
    if pair_train_paths:
        for data_name, data_paths in pair_train_paths.items():
            logger.info(f"添加数据迭代器，data_name:{data_name}, data_paths:{data_paths}")
            pair_data_loaders.append(
                DataLoader(
                    dataset=PairDataSet(data_paths=data_paths, data_name=data_name, model_name=model_name,
                                        pair_label_map=pair_label_map),
                    shuffle=True,
                    collate_fn=lambda x: pair_collate_fn(x, tokenizer, max_length),
                    drop_last=True,
                    batch_size=pair_bsz,
                    num_workers=2
                )
            )

    # 加载模型、tokenizer
    model = MODEL_NAME_INFO[model_name][0].from_pretrained(
        model_dir,
        trust_remote_code=True,
        torch_dtype=torch.float16
    )
    model.to(device)

    original_weight = get_mean_params(model)

    tokenizer = MODEL_NAME_INFO[model_name][1].from_pretrained(model_dir, trust_remote_code=True)
    if grad_checkpoint:
        try:
            model.gradient_checkpointing_enable()
        except:
            logger.error("gradient_checkpointing failed")
    model.train()

    # 开始训练
    args = TrainingArguments(
        **hf_args,
        torch_compile=torch.__version__.startswith("2"),
        prediction_loss_only=True
    )

    if hf_args["gradient_checkpointing"]:
        args.ddp_find_unused_parameters = False
    # save model by call back, do not need official save function
    args.save_strategy = "no"
    trainer = MyTrainer(
        model=model,
        args=args,
        data_collator=lambda x: x[0],
        train_dataset=VecDataSet(in_batch_data_loaders + pair_data_loaders, loader_idxs),
        tokenizer=tokenizer,
        callbacks=[SaveModelCallBack(output_dir=output_dir, save_steps=conf["save_steps"], local_rank=local_rank)]
    )
    trainer.train()

run_train.py
SmartLi8/stella
from .data_collator import (
    InBatchDataSet,
    in_batch_collate_fn,
    PairDataSet,
    pair_collate_fn,
    comb_data_loader,
    VecDataSet
)

from .utils import cosent_loss, get_mean_params, SaveModelCallBack

src/__init__.py
SmartLi8/stella
"""
扩展当前BERT的长度，新扩展的emebdding用层次分解的位置编码进行初始化
"""
import torch
import json
import os
import shutil

if __name__ == "__main__":
    read_dir = r"E:\PublicModels\piccolo-base-zh"
    save_dir = r"E:\PublicModels\piccolo-base-zh-1024"

    ori_pos = 512
    new_pos = 1024
    hp_alpha = 0.2
    # 创建目录
    os.makedirs(save_dir, exist_ok=True)
    # 先拷贝无关文件
    for name in os.listdir(read_dir):
        if name not in ["pytorch_mdoel.bin", "config.json", "tokenizer_config.json"]:
            shutil.copy(os.path.join(read_dir, name), os.path.join(save_dir, name))
    # config.json
    with open(os.path.join(read_dir, "config.json"), "r", encoding="utf8") as fr:
        data = json.load(fr)
    data["max_position_embeddings"] = new_pos
    with open(os.path.join(save_dir, "config.json"), "w", encoding="utf8") as fw:
        data = json.dump(data, fw, ensure_ascii=False, indent=1)

    # tokenizer_config.json
    with open(os.path.join(read_dir, "tokenizer_config.json"), "r", encoding="utf8") as fr:
        data = json.load(fr)
    data["model_max_length"] = new_pos
    with open(os.path.join(save_dir, "tokenizer_config.json"), "w", encoding="utf8") as fw:
        data = json.dump(data, fw, ensure_ascii=False, indent=1)

    # 处理pytorch
    ori_dict = torch.load(os.path.join(read_dir, "pytorch_model.bin"))
    ori_dict["embeddings.position_ids"] = torch.LongTensor([list(range(new_pos))])

    ori_embed = ori_dict["embeddings.position_embeddings.weight"]  # shape [512, 768]
    position_ids = torch.LongTensor(list(range(new_pos)))  # [0,1,2,3,....1023]
    i = position_ids // 512
    j = position_ids % 512
    base_embedding = (ori_embed - ori_embed[0:1] * hp_alpha) / (1 - hp_alpha)

    position_embeddings = hp_alpha * base_embedding[i] + (1 - hp_alpha) * base_embedding[j]
    print("position_embeddings.shape", position_embeddings.shape)
    ori_dict["embeddings.position_embeddings.weight"] = position_embeddings

    torch.save(ori_dict, os.path.join(save_dir, "pytorch_model.bin"))

src/add_new_pos_embed.py
SmartLi8/stella
# coding=utf8
import collections
import random
import torch
from torch.utils.data import Dataset
import json
from loguru import logger


class InBatchDataSet(Dataset):
    def __init__(self, data_paths: str, data_name: str, model_name: str):
        self.data = []
        self.data_paths = data_paths
        self.data_name = data_name
        self.model_name = model_name
        self.load_data()

    def load_data(self):
        for data_path in self.data_paths:
            with open(data_path, "r", encoding="utf8") as fr:
                single_data = [json.loads(line) for line in fr][:]
                self.data.extend(single_data)
        self.data = [[item["txt1"], item["txt2"], item.get("hard_negs", [])] for item in self.data]
        if self.model_name in ["bge", "piccolo"]:
            logger.info(f"检测到是{self.model_name}模型，对于q-p数据前面添加特定的instruction")
            num_added = 0
            # query前面加东西
            for item in self.data:
                txt1, txt2 = item[:2]
                if len(txt1) < 32 and len(txt2) > 64:
                    num_added += 1
                    if self.model_name == "piccolo":
                        item[0] = f"查询: {txt1}"
                        item[1] = f"结果: {txt2}"
                    else:
                        item[0] = f"为这个句子生成表示以用于检索相关文章：{txt1}"
            logger.info(f"数据总量：{len(self.data)}，添加特定指示的数据量：{num_added}")
        self.data = [[self.data_name] + i for i in self.data]

    def __len__(self):
        return len(self.data)

    def __getitem__(self, item):
        """
        item 为数据索引，迭代取第item条数据
        """
        return self.data[item]


def in_batch_collate_fn(batch, tokenizer, max_length):
    """
    DataLoader类collate_fn函数。用于batch内负采样数据处理。
    :param batch: 一个batch内样本数据
    :param tokenizer:tokenizer处理器
    :param max_length:最大长度
    :return:batch内负采样处理后数
    """
    data_name = batch[0][0]
    batch = [item[1:] for item in batch]
    # 随机获取一批难负例
    # hard_negs = []
    # 以hard_neg_ratio的比例进行难负例采样
    hard_negs = [random.choice(negs) for _, _, negs in batch if negs and random.random() < 0.2]
    # print("len(hard_negs)", len(hard_negs))
    batch = [[t1, t2] for t1, t2, _ in batch]

    # q之间不能有重复
    batch = list(dict(batch).items())
    batch = [item[::-1] for item in batch]

    # p之间不能有重复
    batch = list(dict(batch).items())
    batch = [item[::-1] for item in batch]

    # q和p之间不能互相有重复
    new_batch = []
    q_set = collections.Counter([i[0] for i in batch])
    p_set = collections.Counter([i[1] for i in batch])
    for q, p in batch:
        if q != p:
            if q not in p_set and p not in q_set:
                new_batch.append([q, p])
        else:
            new_batch.append([q, p])  # ???
    batch = new_batch

    pos_texts = set([j for item in batch for j in item])
    hard_negs = [i for i in hard_negs if i not in pos_texts]
    all_texts = [item[0] for item in batch] + [item[1] for item in batch] + hard_negs
    ipts = []
    # print("len(all_texts)", len(all_texts))
    for start in range(0, len(all_texts), 32):
        ipt = tokenizer.batch_encode_plus(
            all_texts[start:start + 32], padding="longest", truncation=True, max_length=max_length, return_tensors="pt")
        # print("in_batch_collate_fn", ipt["input_ids"].shape)
        ipts.append(ipt)
    # 最后把q数量加上
    ipts.append(len(batch))
    return [f"in_batch-{data_name}"] + ipts


class PairDataSet(Dataset):
    """ pair对数据集 """

    def __init__(self, data_paths: str, data_name: str, model_name: str, pair_label_map: dict):
        self.data = []
        self.data_paths = data_paths
        self.data_name = data_name
        self.model_name = model_name
        self.pair_label_map = pair_label_map
        self.load_data()

    def load_data(self):
        for data_path in self.data_paths:
            with open(data_path, "r", encoding="utf8") as fr:
                single_data = [json.loads(line) for line in fr]
                self.data.extend(single_data)
        self.data = [
            [item["txt1"], item["txt2"], self.pair_label_map[str(item["label"])]]
            for item in self.data if str(item["label"]) in self.pair_label_map
        ]
        self.data = [[self.data_name] + i for i in self.data]
        logger.info(f"PairDataSet数据总量：{len(self.data)}")

    def __len__(self):
        return len(self.data)

    def __getitem__(self, item):
        """
        item 为数据索引，迭代取第item条数据
        """
        return self.data[item]


def pair_collate_fn(batch, tokenizer, max_length):
    """
    DataLoader类collate_fn函数。用于pair对数据集处理。
    :param batch: List[data_set[i]]
    :param tokenizer:tokenizer处理器
    :param max_length:最大长度
    :return:
    """
    data_name = batch[0][0]
    batch = [item[1:] for item in batch]
    all_texts = [item[0] for item in batch] + [item[1] for item in batch]
    ipts = []
    for start in range(0, len(all_texts), 32):
        ipt = tokenizer.batch_encode_plus(
            all_texts[start:start + 32], padding="longest", truncation=True, max_length=max_length, return_tensors="pt")
        # print("pair_collate_fn", ipt["input_ids"].shape)
        ipts.append(ipt)

    labels = [item[2] for item in batch]
    neg_pos_idxs = [[0.0] * len(batch) for _ in range(len(batch))]
    for i in range(len(batch)):
        for j in range(len(batch)):
            # 1,4: 相似
            # 0 2 3 不相似，其中3是中立，2是完全相反
            if labels[i] in [0, 2, 3] and labels[j] in [1, 4]:
                neg_pos_idxs[i][j] = 1.0
            elif labels[i] == 2 and labels[j] == 3:
                neg_pos_idxs[i][j] = 1.0

    return [f"pair-{data_name}"] + ipts + [torch.tensor(neg_pos_idxs), len(batch)]


def comb_data_loader(loaders, idx_list=None):
    if idx_list is None:
        idx_list = list(range(len(loaders)))
    loaders_iter = [iter(item) for item in loaders]
    idx_for_idx = 0
    while True:
        loader_idx = idx_list[idx_for_idx]
        try:
            yield next(loaders_iter[loader_idx])
        except StopIteration:
            loaders_iter[loader_idx] = iter(loaders[loader_idx])
            yield next(loaders_iter[loader_idx])
        idx_for_idx += 1
        if idx_for_idx % len(idx_list) == 0:
            random.shuffle(idx_list)
            idx_for_idx = 0


class VecDataSet(Dataset):
    """ pair 对数据集 """

    def __init__(self, data_loaders, loader_idxs):
        self.lens = sum([len(i) for i in data_loaders])
        self.data = comb_data_loader(data_loaders, idx_list=loader_idxs)

    def __len__(self):
        return self.lens

    def __getitem__(self, item):
        """
        item 为数据索引，迭代取第item条数据
        """
        return next(self.data)

src/data_collator.py
SmartLi8/stella
""" eval for stella model """
import numpy as np
import torch
import random
import argparse
import functools
from typing import List, Dict
from sentence_transformers import SentenceTransformer
from mteb import MTEB, DRESModel

from C_MTEB.tasks import *

parser = argparse.ArgumentParser(description='evaluation for CMTEB')
parser.add_argument('--model_name', default='bert-base-uncased',
                    type=str, help='which model to use')
parser.add_argument('--output_dir', default='zh_results/',
                    type=str, help='output directory')
parser.add_argument('--max_len', default=512, type=int, help='max length')

args = parser.parse_args()


class RetrievalModel(DRESModel):
    def __init__(self, encoder, **kwargs):
        self.encoder = encoder

    def encode_queries(self, queries: List[str], **kwargs) -> np.ndarray:
        input_texts = ['查询: {}'.format(q) for q in queries]
        return self._do_encode(input_texts)

    def encode_corpus(self, corpus: List[Dict[str, str]], **kwargs) -> np.ndarray:
        input_texts = ['{} {}'.format(doc.get('title', ''), doc['text']).strip() for doc in corpus]
        input_texts = ['结果: {}'.format(t) for t in input_texts]
        return self._do_encode(input_texts)

    @torch.no_grad()
    def _do_encode(self, input_texts: List[str]) -> np.ndarray:
        return self.encoder.encode(
            sentences=input_texts,
            batch_size=256,
            normalize_embeddings=True,
            convert_to_numpy=True
        )


TASKS_WITH_PROMPTS = ["T2Retrieval", "MMarcoRetrieval", "DuRetrieval", "CovidRetrieval", "CmedqaRetrieval",
                      "EcomRetrieval", "MedicalRetrieval", "VideoRetrieval"]

# python run_eval_stellapy  --model_name stella-base-zh  --output_dir ./zh_results/stella-base

if __name__ == '__main__':
    # 加载模型，使用half来加速
    encoder = SentenceTransformer(args.model_name).half()
    encoder.encode = functools.partial(encoder.encode, normalize_embeddings=True)
    encoder.max_seq_length = int(args.max_len)
    # 获取所有任务
    task_names = [t.description["name"] for t in MTEB(task_langs=['zh', 'zh-CN']).tasks]
    random.shuffle(task_names)
    print("task数量", len(task_names))
    print("task_names", task_names)
    for task in task_names:
        evaluation = MTEB(tasks=[task], task_langs=['zh', 'zh-CN'])
        if task in TASKS_WITH_PROMPTS:
            evaluation.run(RetrievalModel(encoder), output_folder=args.output_dir, overwrite_results=False)
        else:
            evaluation.run(encoder, output_folder=args.output_dir, overwrite_results=False)

src/run_eval_stella.py
SmartLi8/stella
# coding=utf8

from transformers import TrainingArguments
from transformers import TrainerCallback, TrainerControl, TrainerState
import torch
from os.path import join


def get_mean_params(model):
    """

    :param model:
    :return:Dict[para_name, para_weight]
    """
    result = {}
    for param_name, param in model.named_parameters():
        result[param_name] = param.data.clone()
    return result


def cosent_loss(neg_pos_idxs, pred_sims, cosent_ratio, zero_data):
    pred_sims = pred_sims * cosent_ratio
    pred_sims = pred_sims[:, None] - pred_sims[None, :]  # 这里是算出所有位置 两两之间余弦的差值
    pred_sims = pred_sims - (1 - neg_pos_idxs) * 1e12
    pred_sims = pred_sims.view(-1)
    pred_sims = torch.cat((zero_data, pred_sims), dim=0)  # 这里加0是因为e^0 = 1相当于在log中加了1
    return torch.logsumexp(pred_sims, dim=0)


class SaveModelCallBack(TrainerCallback):
    def __init__(self, output_dir, save_steps, local_rank):
        self.customized_save_steps = save_steps
        self.customized_output_dir = output_dir
        self.local_rank = local_rank

    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
        if self.local_rank == 0 and self.customized_save_steps > 0 and state.global_step > 0 and state.global_step % self.customized_save_steps == 0:
            epoch = int(state.epoch)
            save_dir = join(self.customized_output_dir, f"epoch-{epoch}_globalStep-{state.global_step}")
            kwargs["model"].save_pretrained(save_dir, max_shard_size="900000MB")
            kwargs["tokenizer"].save_pretrained(save_dir)
            kwargs["tokenizer"].save_vocabulary(save_dir)

src/utils.py
SmartLi8/stella