NLP（四十三）：sentence_bert+pytorch向量检索，进行语义匹配

一、项目目录

 二、data_clean生成数据

from common.root_path import root
import os
import pandas as pd

class DataMerge(object):
    def __init__(self):
        self.data_path = os.path.join(root, "data", "raw_data")
        self.out_path = os.path.join(root, "data", "clean_data")
        self.neg_data = os.path.join(self.out_path, "negtive_data.txt")
        self.sim_path = os.path.join(self.out_path, "sim_data.txt")
        self.final = os.path.join(root, "data", "final_data")
        self.train_sim = os.path.join(self.final, "train.txt")
        self.dev_sim = os.path.join(self.final, "dev.txt")
        self.test_sim = os.path.join(self.final, "test.txt")

    def data_merge(self, role):
        out_sentence, out_label = [], []
        data_path = os.path.join(self.data_path, role)
        # 训练集数据
        train_path = os.path.join(data_path, "train.txt")
        train_t = pd.read_csv(train_path, sep="\t", header=None, names=["sentence", "label"])
        out_sentence.extend(train_t["sentence"])
        out_label.extend(train_t["label"])
        # 验证集数据
        dev_path = os.path.join(data_path, "dev.txt")
        dev_t = pd.read_csv(dev_path, sep="\t", header=None, names=["sentence", "label"])
        out_sentence.extend(dev_t["sentence"])
        out_label.extend(dev_t["label"])
        # 测试集数据
        test_path = os.path.join(data_path, "test.txt")
        test_t = pd.read_csv(test_path, sep="\t", header=None, names=["sentence", "label"])
        out_sentence.extend(test_t["sentence"])
        out_label.extend(test_t["label"])
        # 去重
        clean_sentence, clean_label = [], []
        for s,l in zip(out_sentence, out_label):
            if s not in clean_sentence:
                clean_sentence.append(s)
                clean_label.append(l)
        # 写入文件
        df = pd.DataFrame(
            {
                "sentence": clean_sentence,
                "label": clean_label,
            }
        )
        all_data_path = os.path.join(self.out_path, "all_" + role + ".txt")
        df.to_csv(all_data_path, sep="\t", index=None, header=None)
        return df

    def get_all(self):
        """坐席客户数据文件合并"""
        df_seats = self.data_merge("seats")
        df_semantic = self.data_merge("semantic")
        df = pd.concat([df_seats, df_semantic])
        all_data = os.path.join(self.out_path, "all_data.txt")
        df.to_csv(all_data, sep="\t", header=None, index=None)

    def generator_negtive_data(self):
        """数据随机打乱"""
        all_data = os.path.join(self.out_path, "all_data.txt")
        df_1 = pd.read_csv(all_data, sep="\t", header=None, names=["sentence", "label"])
        df_2 = df_1.sample(frac=1)
        out_s1, out_s2, out_l = [],[],[]
        for s1, l1, s2, l2 in zip(df_1["sentence"], df_1["label"],
                                  df_2["sentence"], df_2["label"]):
            l1 = l1.replace(" ", "").replace("\n", "").replace("\r", "")

            if l1 != l2:
                out_s1.append(s1)
                out_s2.append(s2)
                out_l.append("0")
        df = pd.DataFrame({
            "s1":out_s1,
            "s2":out_s2,
            "label": out_l
        })
        df.to_csv(self.neg_data, sep="\t", index=None)

    def generator_sim_data(self):
        out_s1, out_s2, out_label = list(), list(), list()
        all_data = os.path.join(self.out_path, "all_data.txt")
        t = pd.read_csv(all_data, sep="\t", header=None, names=["sentence", "label"])
        data_dict = dict()
        for index, row in t.iterrows():
            s = row["sentence"]
            l = row["label"]
            if l not in data_dict.keys():
                data_dict[l] = list()
            data_dict[l].append(s)
        for l, s_list in data_dict.items():
            s_list_len = len(s_list)
            for index, s in enumerate(s_list):
                if index > s_list_len -2:
                    break
                out_s1.append(s)
                out_s2.append(s_list[index + 1])
                out_label.append("1")

        df = pd.DataFrame({
            "s1": out_s1,
            "s2": out_s2,
            "label": out_label
        })
        df.to_csv(self.sim_path, index=None, sep="\t")

    def merge_sim_neg_data(self):
        df1 = pd.read_csv(self.sim_path, sep="\t")
        df2 = pd.read_csv(self.neg_data, sep="\t")
        df_all = pd.concat([df1, df2])
        df = df_all.sample(frac=1.0)
        cut_idx_1 = int(round(0.05 * df.shape[0]))
        cut_idx_2 = int(round(0.1 * df.shape[0]))
        print(cut_idx_1, cut_idx_2)
        df_test, df_dev, df_train = df.iloc[:cut_idx_1], df.iloc[cut_idx_1:cut_idx_2], df.iloc[cut_idx_2:]
        df_test.to_csv(self.test_sim, index=False, sep='\t')
        df_dev.to_csv(self.dev_sim, index=False, sep='\t')
        df_train.to_csv(self.train_sim, index=False, sep='\t')

if __name__ == '__main__':
    DataMerge().merge_sim_neg_data()

三、root_path

import os
__all__ = ["root"]
_parent_path = os.path.split(os.path.realpath(__file__))[0]
_root = _parent_path[:_parent_path.find("sentence_bert")]
root = os.path.join(_root, "sentence_bert")

四、训练

from torch.utils.data import DataLoader
import math
from sentence_transformers import SentenceTransformer,  LoggingHandler, losses, models, util
from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator
from sentence_transformers.readers import InputExample
import logging
from datetime import datetime
import os
from common.root_path import root
import pandas as pd

class MySentenceBert():
    logging.basicConfig(format='%(asctime)s - %(message)s',
                        datefmt='%Y-%m-%d %H:%M:%S',
                        level=logging.INFO,
                        handlers=[LoggingHandler()])
    def __init__(self):

        self.train_batch_size = 16
        self.num_epochs = 4
        data_path = os.path.join(root, "data", "final_data")
        self.train_data = pd.read_csv(os.path.join(data_path, "train.txt"), sep="\t")
        self.val_data = pd.read_csv(os.path.join(data_path, "val.txt"), sep="\t")
        self.test_data = pd.read_csv(os.path.join(data_path, "test.txt"), sep="\t")
        self.model_save_path = os.path.join(root, "chkpt", "sentence_bert_model" +
                                            datetime.now().strftime("_%Y_%m_%d_%H_%M"))

    def data_generator(self):
        logging.info("generator dataset")
        train_datas = []
        dev_datas = []
        test_datas = []
        for s1, s2, l in zip(self.train_data["s1"],
                             self.train_data["s2"],
                             self.train_data["label"]):
            train_datas.append(InputExample(texts=[s1, s2], label=float(l)))
        for s1, s2, l in zip(self.val_data["s1"],
                             self.val_data["s2"],
                             self.val_data["label"]):
            dev_datas.append(InputExample(texts=[s1, s2], label=float(l)))
        for s1, s2, l in zip(self.test_data["s1"],
                             self.test_data["s2"],
                             self.test_data["label"]):
            test_datas.append(InputExample(texts=[s1, s2], label=float(l)))
        return train_datas, dev_datas, test_datas

    def train(self, train_datas, dev_datas, model):
        train_dataloader = DataLoader(train_datas, shuffle=True, batch_size=self.train_batch_size)
        train_loss = losses.CosineSimilarityLoss(model=model)
        evaluator = EmbeddingSimilarityEvaluator.from_input_examples(dev_datas, name='sts-dev')
        warmup_steps = math.ceil(len(train_dataloader) * self.num_epochs  * 0.1)
        logging.info("Warmup-steps: {}".format(warmup_steps))
        model.fit(train_objectives=[(train_dataloader, train_loss)],
                  evaluator=evaluator,
                  epochs=self.num_epochs,
                  evaluation_steps=1000,
                  warmup_steps=warmup_steps,
                  output_path=self.model_save_path)
    def test(self, test_samples):
        model = SentenceTransformer(self.model_save_path)
        test_evaluator = EmbeddingSimilarityEvaluator.from_input_examples(test_samples, name='sts-test')
        test_evaluator(model, output_path=self.model_save_path)

    def main(self):
        train_datas, dev_datas, test_datas = self.data_generator()

        model_name = os.path.join(root, "chkpt", "distiluse-base-multilingual-cased")
        word_embedding_model = models.Transformer(model_name)
        pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(),
                                       pooling_mode_mean_tokens=True,
                                       pooling_mode_cls_token=False,
                                       pooling_mode_max_tokens=False)
        model = SentenceTransformer(modules=[word_embedding_model, pooling_model])
        self.train(train_datas, dev_datas, model)
        self.test(test_datas)

if __name__ == '__main__':
    MySentenceBert().main()

五、向量检索

from sentence_transformers import SentenceTransformer, util
import os
import csv
import pickle
import time
from root_path import root
import json

class SemanticSearch():
    def __init__(self):
        model_name = os.path.join(root, "chkpt", "sentence_bert_model_2021_08_05_18_16")
        self.model = SentenceTransformer(model_name)
        embedding_cache_path = 'semantic_search_embedding.pkl'
        dataset_path = os.path.join(root, "data", "bert_data", "index.txt")
        with open(os.path.join(root, "config", "code_to_label.json"), "r", encoding="utf8") as f:
            self.d = json.load(f)
        self.sentences = list()
        self.code = list()
        if not os.path.exists(embedding_cache_path):
            with open(dataset_path, encoding='utf8') as fIn:
                for read_line in fIn:
                    read_line = read_line.split("\t")
                    self.sentences.append(read_line[0])
                    self.code.append(read_line[1].replace("\n", ""))
            print("Encode the corpus. This might take a while")
            self.embeddings = self.model.encode(self.sentences, show_progress_bar=True, convert_to_tensor=True)
            print("Store file on disc")
            with open(embedding_cache_path, "wb") as fOut:
                pickle.dump({'sentences': self.sentences, 'embeddings': self.embeddings, "code": self.code}, fOut)
        else:
            print("Load pre-computed embeddings from disc")
            with open(embedding_cache_path, "rb") as fIn:
                cache_data = pickle.load(fIn)
                self.sentences = cache_data['sentences']
                self.embeddings = cache_data['embeddings']
                self.code = cache_data["code"]

    def query(self, query):
        inp_question = query
        question_embedding = self.model.encode(inp_question, convert_to_tensor=True)
        hits = util.semantic_search(question_embedding, self.embeddings)
        hit = hits[0][0]  # Get the hits for the first query
        score = hit['score']
        text = self.sentences[hit['corpus_id']]
        kh_code = self.code[hit['corpus_id']]
        label = self.d[kh_code][1]
        return label,score,text

    def main(self):
        self.query("你好")


if __name__ == '__main__':
    SemanticSearch().main()

 六、参考

https://github.com/UKPLab/sentence-transformers/blob/master/examples/training/sts/training_stsbenchmark.py

https://github.com/UKPLab/sentence-transformers/blob/master/examples/applications/semantic-search/semantic_search_quora_pytorch.py