语义相似度与 BERT

import numpy as np
import pandas as pd
import tensorflow as tf
import transformers

max_length = 128  # 输入句子的最大长度
batch_size = 32
epochs = 2

# 我们数据集中的标签。
labels = ["contradiction", "entailment", "neutral"]

!curl -LO https://raw.githubusercontent.com/MohamadMerchant/SNLI/master/data.tar.gz
!tar -xvzf data.tar.gz

# 总共有超过 55 万个样本；我们将使用 10 万个作为此示例。
train_df = pd.read_csv("SNLI_Corpus/snli_1.0_train.csv", nrows=100000)
valid_df = pd.read_csv("SNLI_Corpus/snli_1.0_dev.csv")
test_df = pd.read_csv("SNLI_Corpus/snli_1.0_test.csv")

# 数据的形状
print(f"总训练样本 : {train_df.shape[0]}")
print(f"总验证样本: {valid_df.shape[0]}")
print(f"总测试样本: {valid_df.shape[0]}")

print(f"句子1: {train_df.loc[1, 'sentence1']}")
print(f"句子2: {train_df.loc[1, 'sentence2']}")
print(f"相似度: {train_df.loc[1, 'similarity']}")

# 我们的训练数据中有一些 NaN 条目，我们将简单地删除它们。
print("缺失值的数量")
print(train_df.isnull().sum())
train_df.dropna(axis=0, inplace=True)

print("训练目标分布")
print(train_df.similarity.value_counts())

print("验证目标分布")
print(valid_df.similarity.value_counts())

train_df = (
    train_df[train_df.similarity != "-"]
    .sample(frac=1.0, random_state=42)
    .reset_index(drop=True)
)
valid_df = (
    valid_df[valid_df.similarity != "-"]
    .sample(frac=1.0, random_state=42)
    .reset_index(drop=True)
)

train_df["label"] = train_df["similarity"].apply(
    lambda x: 0 if x == "contradiction" else 1 if x == "entailment" else 2
)
y_train = tf.keras.utils.to_categorical(train_df.label, num_classes=3)

valid_df["label"] = valid_df["similarity"].apply(
    lambda x: 0 if x == "contradiction" else 1 if x == "entailment" else 2
)
y_val = tf.keras.utils.to_categorical(valid_df.label, num_classes=3)

test_df["label"] = test_df["similarity"].apply(
    lambda x: 0 if x == "contradiction" else 1 if x == "entailment" else 2
)
y_test = tf.keras.utils.to_categorical(test_df.label, num_classes=3)

class BertSemanticDataGenerator(tf.keras.utils.Sequence):
    """生成数据批次。

    参数：
        sentence_pairs: 前提和假设输入句子的数组。
        labels: 标签数组。
        batch_size: 整数批量大小。
        shuffle: 布尔值，是否打乱数据。
        include_targets: 布尔值，是否包括标签。

    返回：
        元组 `([input_ids, attention_mask, `token_type_ids], labels)`
        （或者仅仅是 `[input_ids, attention_mask, `token_type_ids]`
         如果 `include_targets=False`）
    """

    def __init__(
        self,
        sentence_pairs,
        labels,
        batch_size=batch_size,
        shuffle=True,
        include_targets=True,
    ):
        self.sentence_pairs = sentence_pairs
        self.labels = labels
        self.shuffle = shuffle
        self.batch_size = batch_size
        self.include_targets = include_targets
        # 加载我们的BERT分词器以编码文本。
        # 我们将使用基础的无大小写预训练模型。
        self.tokenizer = transformers.BertTokenizer.from_pretrained(
            "bert-base-uncased", do_lower_case=True
        )
        self.indexes = np.arange(len(self.sentence_pairs))
        self.on_epoch_end()

    def __len__(self):
        # 表示每个周期的批次数量。
        return len(self.sentence_pairs) // self.batch_size

    def __getitem__(self, idx):
        # 检索索引的批次。
        indexes = self.indexes[idx * self.batch_size : (idx + 1) * self.batch_size]
        sentence_pairs = self.sentence_pairs[indexes]

        # 使用BERT分词器的batch_encode_plus函数将两个句子的批量
        # 一起编码并通过[SEP]标记分隔。
        encoded = self.tokenizer.batch_encode_plus(
            sentence_pairs.tolist(),
            add_special_tokens=True,
            max_length=max_length,
            return_attention_mask=True,
            return_token_type_ids=True,
            pad_to_max_length=True,
            return_tensors="tf",
        )

        # 将编码特征的批量转换为numpy数组。
        input_ids = np.array(encoded["input_ids"], dtype="int32")
        attention_masks = np.array(encoded["attention_mask"], dtype="int32")
        token_type_ids = np.array(encoded["token_type_ids"], dtype="int32")

        # 如果数据生成器用于训练/验证，则设为true。
        if self.include_targets:
            labels = np.array(self.labels[indexes], dtype="int32")
            return [input_ids, attention_masks, token_type_ids], labels
        else:
            return [input_ids, attention_masks, token_type_ids]

    def on_epoch_end(self):
        # 如果shuffle被设置为True，则在每个周期结束后打乱索引。
        if self.shuffle:
            np.random.RandomState(42).shuffle(self.indexes)

# 在分布式策略范围内创建模型。
strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
    # 来自BERT分词器的编码token ids。
    input_ids = tf.keras.layers.Input(
        shape=(max_length,), dtype=tf.int32, name="input_ids"
    )
    # 注意力掩码指示模型应该关注哪些tokens。
    attention_masks = tf.keras.layers.Input(
        shape=(max_length,), dtype=tf.int32, name="attention_masks"
    )
    # token类型ids是标识模型中不同序列的二进制掩码。
    token_type_ids = tf.keras.layers.Input(
        shape=(max_length,), dtype=tf.int32, name="token_type_ids"
    )
    # 加载预训练的BERT模型。
    bert_model = transformers.TFBertModel.from_pretrained("bert-base-uncased")
    # 冻结BERT模型以重新使用预训练特征而不进行修改。
    bert_model.trainable = False

    bert_output = bert_model.bert(
        input_ids, attention_mask=attention_masks, token_type_ids=token_type_ids
    )
    sequence_output = bert_output.last_hidden_state
    pooled_output = bert_output.pooler_output
    # 在冻结层上添加可训练层，以便将预训练特征适应新数据。
    bi_lstm = tf.keras.layers.Bidirectional(
        tf.keras.layers.LSTM(64, return_sequences=True)
    )(sequence_output)
    # 对bi_lstm序列输出应用混合池化方法。
    avg_pool = tf.keras.layers.GlobalAveragePooling1D()(bi_lstm)
    max_pool = tf.keras.layers.GlobalMaxPooling1D()(bi_lstm)
    concat = tf.keras.layers.concatenate([avg_pool, max_pool])
    dropout = tf.keras.layers.Dropout(0.3)(concat)
    output = tf.keras.layers.Dense(3, activation="softmax")(dropout)
    model = tf.keras.models.Model(
        inputs=[input_ids, attention_masks, token_type_ids], outputs=output
    )

    model.compile(
        optimizer=tf.keras.optimizers.Adam(),
        loss="categorical_crossentropy",
        metrics=["acc"],
    )


print(f"策略: {strategy}")
model.summary()

train_data = BertSemanticDataGenerator(
    train_df[["sentence1", "sentence2"]].values.astype("str"),
    y_train,
    batch_size=batch_size,
    shuffle=True,
)
valid_data = BertSemanticDataGenerator(
    valid_df[["sentence1", "sentence2"]].values.astype("str"),
    y_val,
    batch_size=batch_size,
    shuffle=False,
)

history = model.fit(
    train_data,
    validation_data=valid_data,
    epochs=epochs,
    use_multiprocessing=True,
    workers=-1,
)

# 解除 bert_model 的冻结。
bert_model.trainable = True
# 重新编译模型以使更改生效。
model.compile(
    optimizer=tf.keras.optimizers.Adam(1e-5),
    loss="categorical_crossentropy",
    metrics=["accuracy"],
)
model.summary()

历史 = 模型.fit(
    训练数据,
    验证数据=验证数据,
    轮数=轮数,
    使用多进程=True,
    工作者=-1,
)

测试数据 = BertSemanticDataGenerator(
    测试_df[["句子1", "句子2"]].值.astype("str"),
    y_test,
    batch_size=batch_size,
    shuffle=False,
)
模型.evaluate(测试数据, verbose=1)

def check_similarity(句子1, 句子2):
    句子对 = np.array([[str(句子1), str(句子2)]])
    测试数据 = BertSemanticDataGenerator(
        句子对, 标签=None, batch_size=1, shuffle=False, include_targets=False,
    )

    概率 = 模型.predict(测试数据[0])[0]
    idx = np.argmax(概率)
    概率 = f"{概率[idx]: .2f}%"
    预测 = 标签[idx]
    return 预测, 概率

句子1 = "两名女性正在一起观察某些东西。"
句子2 = "两名女性正站着，闭着眼睛。"
check_similarity(句子1, 句子2)

句子1 = "一位微笑的穿着戏服的女性正在拿着一把伞"
句子2 = "一位快乐的女性穿着仙女服装，手中拿着一把伞"
check_similarity(句子1, 句子2)

sentence1 = "一个有多个男性参与的足球比赛"
sentence2 = "一些男性正在进行一项运动"
check_similarity(sentence1, sentence2)