KL 散度(python+nlp)

python demo

KL 散度（Kullback-Leibler divergence），也称为相对熵，是衡量两个概率分布之间差异的一种方式。KL 散度是非对称的，也就是说，P 相对于 Q 的 KL 散度通常不等于 Q 相对于 P 的 KL 散度。

一个简单的 Python 类来计算两个离散概率分布之间的 KL 散度：

import numpy as np

class KLDivergence:
    def __init__(self, eps=1e-10):
        self.eps = eps  # 防止出现 log(0)

    def kl_divergence(self, p, q):
        """
        计算两个离散概率分布 P 和 Q 之间的 KL 散度。
        
        参数:
        p (np.array): 分布 P 的概率值。
        q (np.array): 分布 Q 的概率值。
        
        返回:
        float: P 相对于 Q 的 KL 散度。
        """
        p = np.asarray(p, dtype=np.float)
        q = np.asarray(q, dtype=np.float)
        
        # 防止分母为零
        q = np.clip(q, self.eps, 1 - self.eps)
        
        # 计算 KL 散度
        kl_div = np.sum(np.where(p != 0, p * np.log(p / q), 0))
        return kl_div

# 示例代码
if __name__ == "__main__":
    kld = KLDivergence()
    
    # 定义两个概率分布
    p = np.array([0.1, 0.4, 0.5])
    q = np.array([0.2, 0.3, 0.5])
    
    # 计算 KL 散度
    kl_div = kld.kl_divergence(p, q)
    print("KL Divergence:", kl_div)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

nlp demo 1

在自然语言处理（NLP）中，KL 散度可以用于多种场景，比如评估文档的主题分布一致性、语料库中的词频分布比较等。这里提供一个使用 KL 散度来比较两个文档中词频分布的例子。我们将使用 scikit-learn 库来提取文档的词频，并计算它们之间的 KL 散度。

示例说明

在这个例子中，我们将使用 TF-IDF（Term Frequency-Inverse Document Frequency）来表示文档中的词频，并计算两个文档之间的 KL 散度。TF-IDF 是一种常用的文本特征表示方法，它可以衡量一个词对文档的重要程度。

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from scipy.stats import entropy

class DocumentComparator:
    def __init__(self):
        self.vectorizer = TfidfVectorizer()
        self.eps = 1e-10

    def fit_transform(self, documents):
        """
        使用 TF-IDF 向量化器将文档转换为 TF-IDF 特征向量。
        
        参数:
        documents (list of str): 文档列表。
        
        返回:
        np.array: TF-IDF 特征矩阵。
        """
        tfidf_matrix = self.vectorizer.fit_transform(documents)
        return tfidf_matrix.toarray()

    def calculate_kl_divergence(self, doc1_tfidf, doc2_tfidf):
        """
        计算两个文档 TF-IDF 特征向量之间的 KL 散度。
        
        参数:
        doc1_tfidf (np.array): 第一个文档的 TF-IDF 特征向量。
        doc2_tfidf (np.array): 第二个文档的 TF-IDF 特征向量。
        
        返回:
        float: 两个文档之间的 KL 散度。
        """
        # 将 TF-IDF 向量归一化为概率分布
        doc1_prob = doc1_tfidf / np.sum(doc1_tfidf)
        doc2_prob = doc2_tfidf / np.sum(doc2_tfidf)
        
        # 防止分母为零
        doc2_prob = np.clip(doc2_prob, self.eps, 1 - self.eps)
        
        # 计算 KL 散度
        kl_div = entropy(doc1_prob, doc2_prob)
        return kl_div

# 示例代码
if __name__ == "__main__":
    comparator = DocumentComparator()
    
    # 定义两个文档
    doc1 = "Python is a widely used high-level programming language."
    doc2 = "Python is a popular scripting language for data science."
    
    # 将文档转换为 TF-IDF 特征向量
    docs = [doc1, doc2]
    tfidf_matrix = comparator.fit_transform(docs)
    
    # 计算两个文档之间的 KL 散度
    kl_div = comparator.calculate_kl_divergence(tfidf_matrix[0], tfidf_matrix[1])
    print("KL Divergence between documents:", kl_div)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

nlp demo 2

结合 Transformer 模型（如 BERT、GPT-2、T5 或 Llama2）与 KL 散度，我们可以构建一个更复杂的系统，用于评估不同文档或句子之间的相似性。这可以通过以下几种方式实现：

• 使用预训练模型生成句子嵌入：使用 Transformer 模型来生成句子或文档级别的嵌入向量。
• 计算句子嵌入的概率分布：将句子嵌入转换为概率分布。
• 计算 KL 散度：使用 KL 散度来比较两个句子的概率分布。

这里我们使用 Hugging Face 的 Transformers 库来加载预训练的模型，并计算句子之间的 KL 散度。我们将使用 BERT 作为示例模型，但这种方法同样适用于 GPT-2、T5 或 Llama2 等其他 Transformer 模型。

import torch
from transformers import AutoTokenizer, AutoModel
from scipy.special import rel_entr
import numpy as np

class SentenceEmbeddingComparator:
    def __init__(self, model_name="bert-base-uncased"):
        """
        初始化 SentenceEmbeddingComparator 类。
        
        参数:
        model_name (str): 预训练模型的名字，默认为 'bert-base-uncased'。
        """
        # 设置设备为 GPU 如果可用，否则使用 CPU
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # 加载预训练的分词器和模型
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name).to(self.device)
        
        # 设置模型为评估模式
        self.model.eval()

    def _mean_pooling(self, model_output, attention_mask):
        """
        对模型输出执行平均池化以获取句子级别的嵌入。
        
        参数:
        model_output (torch.Tensor): 模型的输出。
        attention_mask (torch.Tensor): 注意力掩码，指示哪些位置是填充的。
        
        返回:
        torch.Tensor: 句子级别的嵌入。
        """
        # 获取最后一层的隐藏状态
        token_embeddings = model_output.last_hidden_state
        
        # 扩展注意力掩码以匹配嵌入的维度
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        
        # 计算每句话的嵌入的加权和
        sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
        
        # 计算每个位置的有效掩码数量
        sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
        
        # 计算平均嵌入
        return sum_embeddings / sum_mask

    def generate_sentence_embedding(self, sentence):
        """
        生成给定句子的嵌入。
        
        参数:
        sentence (str): 输入的句子。
        
        返回:
        np.array: 句子的嵌入向量。
        """
        # 对输入句子进行编码并添加到设备上
        encoded_input = self.tokenizer(sentence, padding=True, truncation=True, max_length=128, return_tensors='pt').to(self.device)
        
        # 使用模型生成输出
        with torch.no_grad():
            model_output = self.model(**encoded_input)
        
        # 使用平均池化获取句子嵌入
        sentence_embedding = self._mean_pooling(model_output, encoded_input['attention_mask'])
        
        # 将嵌入从张量转换为 NumPy 数组
        return sentence_embedding.cpu().numpy()[0]

    def calculate_kl_divergence(self, emb1, emb2):
        """
        计算两个句子嵌入之间的 KL 散度。
        
        参数:
        emb1 (np.array): 第一个句子的嵌入。
        emb2 (np.array): 第二个句子的嵌入。
        
        返回:
        float: 两个句子之间的 KL 散度。
        """
        # 将嵌入向量转换为概率分布
        # 归一化确保概率分布的总和为 1
        prob1 = emb1 / np.linalg.norm(emb1, ord=1)
        prob2 = emb2 / np.linalg.norm(emb2, ord=1)
        
        # 计算 KL 散度
        kl_div = np.sum(rel_entr(prob1, prob2))
        return kl_div

# 示例代码
if __name__ == "__main__":
    # 创建 SentenceEmbeddingComparator 实例
    comparator = SentenceEmbeddingComparator("bert-base-uncased")
    
    # 定义两个句子
    sentence1 = "Python is a widely used high-level programming language."
    sentence2 = "Python is a popular scripting language for data science."
    
    # 生成句子嵌入
    emb1 = comparator.generate_sentence_embedding(sentence1)
    emb2 = comparator.generate_sentence_embedding(sentence2)
    
    # 计算两个句子之间的 KL 散度
    kl_div = comparator.calculate_kl_divergence(emb1, emb2)
    
    # 输出 KL 散度
    print("KL Divergence between sentences:", kl_div)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110