import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import matplotlib.pyplot as pltclass BiLSTM_Attention(nn.Module):
def __init__(self):
super(BiLSTM_Attention, self).__init__()
# 嵌入层，用于将文本中的词汇映射为密集向量表示
# nn.Embedding 是PyTorch提供的用于词嵌入（Word Embedding）的层。
# vocab_size 是词汇表大小，代表有多少个不同的词汇。
# embedding_dim 是嵌入向量的维度，它将每个词汇映射到一个具有 embedding_dim 维度的向量空间中
self.embedding = nn.Embedding(vocab_size, embedding_dim)
# 双向LSTM层，处理嵌入后的词向量，生成LSTM输出
# nn.LSTM 是PyTorch提供的LSTM层。
# embedding_dim 是嵌入向量的维度，也是LSTM层的输入尺寸（input_size），代表每个时间步的输入特征维度。
# n_hidden 是LSTM层的隐藏单元数（hidden_size），代表每个时间步的隐藏状态的维度
# bidirectional=True 表示该LSTM是双向的，即同时考虑正向和反向的序列信息，增加了模型对序列信息的理解能力。
# 这个LSTM层将用于处理输入的嵌入向量，生成LSTM的输出和最终的隐藏状态
self.lstm = nn.LSTM(embedding_dim, n_hidden, bidirectional=True)
# 双向LSTM层，处理嵌入后的词向量，生成LSTM输出
# nn.Linear 是PyTorch提供的全连接（线性）层。
# n_hidden * 2 表示输入特征的维度，由于使用了双向LSTM，所以将正向和反向的隐藏状态拼接在一起，维度变为原来的两倍。
# num_classes 表示分类问题的类别数，因为这是一个二分类任务，所以设为2。
# 这个输出层将用于将LSTM的输出转换为最终的分类结果。输出为2维，其中一个维度表示负面情感的概率，另一个维度表示正面情感的概率
self.out = nn.Linear(n_hidden * 2, num_classes)
# 该函数用于计算注意力权重，然后根据注意力权重对LSTM输出进行加权平均得到上下文向量
# lstm_output : LSTM的输出，维度为 [batch_size, n_step, n_hidden * num_directions(=2)]
# final_state : LSTM最终的隐藏状态，维度为 [num_layers(=1) * num_directions(=2), batch_size, n_hidden]
def attention_net(self, lstm_output, final_state):
# final_state 是LSTM在输入序列最后一个时间步的隐藏状态，维度为 [num_layers(=1) * num_directions(=2), batch_size, n_hidden]
# n_hidden * 2 表示双向LSTM的隐藏状态的维度，乘以2是因为它包含了前向和后向两个方向的隐藏状态。
# final_state.view(-1, n_hidden * 2, 1) 将 final_state 的维度从 [num_layers * num_directions, batch_size, n_hidden]
# 转换为 [batch_size, n_hidden * 2, 1]，在第三维度上增加了一个维度，这是为了后续计算注意力权重做准备
# hidden : [batch_size, n_hidden * num_directions(=2), 1(=n_layer)]
hidden = final_state.view(-1, n_hidden * 2, 1)
# lstm_output 是LSTM的输出，维度为 [batch_size, n_step, n_hidden * num_directions(=2)]。
# hidden 是通过上面的操作得到的 final_state 的改变维度后的表示，维度为 [batch_size, n_hidden * 2, 1]。
# torch.bmm(lstm_output, hidden) 进行矩阵相乘，实现对 lstm_output 和 hidden 进行注意力权重的计算。注意力权重表示LSTM输出中每个时间步对应的重要程度。
# squeeze(2) 将注意力权重张量的第三个维度（为1维）挤压掉，得到维度为 [batch_size, n_step] 的 attn_weights
attn_weights = torch.bmm(lstm_output, hidden).squeeze(2) # attn_weights : [batch_size, n_step]
# F.softmax(attn_weights, 1) 对注意力权重 attn_weights 进行softmax操作，将注意力权重转换为概率分布，使得每个时间步的权重值在0到1之间且和为1
soft_attn_weights = F.softmax(attn_weights, 1)
# [batch_size, n_hidden * num_directions(=2), n_step] * [batch_size, n_step, 1] = [batch_size, n_hidden * num_directions(=2), 1]
# lstm_output.transpose(1, 2) 对LSTM输出进行转置，维度从 [batch_size, n_step, n_hidden * num_directions] 变为 [batch_size, n_hidden * num_directions, n_step]。
# soft_attn_weights.unsqueeze(2) 将 soft_attn_weights 张量的维度从 [batch_size, n_step] 扩展为 [batch_size, n_step, 1]，以便进行矩阵相乘。
# torch.bmm(lstm_output.transpose(1, 2), soft_attn_weights.unsqueeze(2)) 实现LSTM输出和注意力权重的加权平均，得到上下文向量。
# squeeze(2) 将上下文向量张量的第三个维度（为1维）挤压掉，得到维度为 [batch_size, n_hidden * num_directions] 的 context
context = torch.bmm(lstm_output.transpose(1, 2), soft_attn_weights.unsqueeze(2)).squeeze(2)
# 返回上下文向量 context 和注意力权重 soft_attn_weights。
# context 是根据注意力权重加权平均得到的LSTM输出的上下文向量，用于表示输入序列的重要信息。
# soft_attn_weights.data.numpy() 将注意力权重 soft_attn_weights 转换为NumPy数组形式并返回，以便后续可视化注意力权重矩阵
return context, soft_attn_weights.data.numpy() # context : [batch_size, n_hidden * num_directions(=2)]
# 前向传播函数
# X : 输入的文本序列，维度为 [batch_size, len_seq]
def forward(self, X):
# 首先将输入进行嵌入操作，然后交换维度以适应LSTM的输入格式
# 这里使用了 nn.Embedding 层，将输入的文本序列 X 中的每个词汇（对应的索引）映射为一个密集向量表示。
# input 变量维度为 [batch_size, len_seq, embedding_dim]
# 其中 batch_size 表示批次大小，len_seq 表示每个句子中的词汇数，embedding_dim 表示词向量的维度。
input = self.embedding(X)
# 交换维度以适应LSTM的输入格式
# LSTM模型的输入需要将序列长度维度放在第一位，因此这里使用 permute 方法交换了 input 的维度，使其变为 [len_seq, batch_size, embedding_dim]
input = input.permute(1, 0, 2)
# 初始化LSTM的隐藏状态和记忆单元状态为全零张量
# 这里创建了两个零张量 hidden_state 和 cell_state
# 用于存储LSTM的初始隐藏状态和记忆单元状态。
# 这里 1*2 表示LSTM的层数乘以双向LSTM的方向数（正向和反向）
# 维度为：[num_layers(=1) * num_directions(=2), batch_size, n_hidden]
hidden_state = torch.zeros(1*2, len(X), n_hidden)
# 维度为：[num_layers(=1) * num_directions(=2), batch_size, n_hidden]
cell_state = torch.zeros(1*2, len(X), n_hidden)
# final_hidden_state, final_cell_state : [num_layers(=1) * num_directions(=2), batch_size, n_hidden]
# 将嵌入后的序列输入LSTM，得到输出和最终的隐藏状态
# 这里调用了 nn.LSTM 层 self.lstm 来进行前向传播。
# output 是LSTM在所有时间步的输出，维度为 [len_seq, batch_size, n_hidden * num_directions(=2)]。
# final_hidden_state 和 final_cell_state 是LSTM的最终隐藏状态和记忆单元状态，维度均为 [num_layers(=1) * num_directions(=2), batch_size, n_hidden]
output, (final_hidden_state, final_cell_state) = self.lstm(input, (hidden_state, cell_state))
# 这里再次使用 permute 方法调整输出维度，使output恢复为 [batch_size, len_seq, n_hidden]
output = output.permute(1, 0, 2)
# 调用 attention_net 函数得到注意力加权的上下文向量。
# 将上下文向量输入输出层，得到最终的分类结果和注意力权重
# 这里调用了在 BiLSTM_Attention 类中定义的 attention_net 方法
# 该方法用于计算注意力权重并将其应用于LSTM输出 output，得到上下文向量 attn_output 和注意力权重矩阵 attention
attn_output, attention = self.attention_net(output, final_hidden_state)
# 将上下文向量输入输出层，得到最终的分类结果和注意力权重
# 这里将上下文向量 attn_output 输入全连接输出层 self.out，得到最终的分类结果(model)，维度为 [batch_size, num_classes]。
# 返回注意力权重矩阵(attention)，维度为 [batch_size, n_step]。
# 注意力权重矩阵表示模型在分类时对文本序列中每个词的关注程度，可以用于可视化和分析模型的注意力行为
return self.out(attn_output), attention
if __name__ == '__main__':
# embedding_dim 是嵌入维度，用于表示文本中的每个词汇的密集向量表示。
# 在自然语言处理任务中，文本数据往往是由离散的词汇组成，而神经网络很难直接处理这种离散形式的数据。
# 因此，需要将文本中的词汇映射为连续的向量表示，以便神经网络能够处理。
# 嵌入层（nn.Embedding）就是用来完成这个映射过程的。
# 它接受一个词汇表的大小（vocab_size）和一个嵌入维度（embedding_dim）作为输入，
# 然后根据词汇表的大小创建一个随机初始化的嵌入矩阵，其中每个词汇对应一个向量，
# 向量的维度为 embedding_dim。模型在训练过程中，会根据文本数据对这些嵌入向量进行学习，使得相似的词汇在嵌入空间中距离更近，有更好的表示能力
embedding_dim = 2
# n_hidden 表示LSTM隐藏层中的隐藏单元数（也称为隐藏状态的维度）。
# 在双向LSTM模型中，LSTM层会有两个方向的隐藏状态，所以总的隐藏单元数会是 n_hidden * 2。
# 在LSTM中，隐藏状态（hidden state）用于存储过去时间步的信息，通过更新和传递隐藏状态，LSTM能够在处理时间序列数据时更好地捕捉长期依赖关系。
# n_hidden 的值会影响LSTM模型的表示能力和学习能力，过小的值可能会导致模型拟合不足，无法捕获数据的复杂模式，而过大的值可能会增加模型复杂性，使得训练过程变得困难。
# 通常，对于不同的任务和数据集，合适的 n_hidden 取值会有所差异。
# 选择合适的 n_hidden 取决于数据集的大小和复杂度，以及任务的复杂性。
# 常见的做法是通过尝试不同的 n_hidden 值，然后根据在验证集上的性能选择最优的值
n_hidden = 5
# num_classes 在这个代码中用于指定文本分类任务的类别数量。在该代码中，文本分类任务有两个类别，即"好"和"不好"，分别用 1 和 0 表示。
# 在模型的输出层，我们使用 self.out = nn.Linear(n_hidden * 2, num_classes) 这一行代码来定义输出层。
# 这个输出层是一个全连接层，将 LSTM 输出的特征（n_hidden * 2 维）转换成最终的分类结果（num_classes 维）。
# 对于本例中的二分类任务，num_classes 为 2，因为我们需要输出两个值（0或1）来表示分类结果。
# 这样，模型的输出维度就对应着两个类别的概率分布。
# 在训练阶段，模型输出的结果会通过 softmax 函数处理，得到对应类别的概率分布。
# 例如，如果模型输出 [0.7, 0.3]，则表示模型认为该文本属于"好"的概率为 0.7，属于"不好"的概率为 0.3。最终分类结果会根据概率值选择概率较大的类别，例如这里会选择"好"类别(0.7 > 0.3)。
# 因此，num_classes 在代码中用于定义输出层的维度，确保模型能够输出正确的类别概率分布，并根据分类结果进行准确的文本分类
num_classes = 2
# 3词汇句子
# sentences = ["我 爱 你", "他 爱 我", "她 喜欢 篮球", "她 喜欢 他", "我 讨厌 你","他 讨厌 我", "我 对不起 你"]
# labels = [1, 1, 1, 1, 0, 0, 0] # 1 是 好, 0 是 不好. 构建词汇表，并为每个词汇赋予一个唯一的索引
sentences = ["我 爱 你", "他 爱 我", "她 喜欢 篮球", "她 喜欢 他", "我 讨厌 你", "他 讨厌 我", "我 对不起 你",
"我 很高兴","他 很高兴","你 很高兴","我 不高兴","他 不高兴","她 不高兴"]
labels = [1, 1, 1, 1, 0, 0, 0, 1,1,1,0,0,0]
# Define the maximum sequence length
max_seq_length = 6
# 为每个句子添加一个特殊的标记（例如<PAD>），使其与最长句子长度保持一致，以便构成一个批次的张量输入
def pad_sequence(sentence, max_len):
words = sentence.split()
if len(words) < max_len:
words.extend(['<PAD>'] * (max_len - len(words)))
return " ".join(words[:max_len])
# 对所有句子进行填充
sentences = [pad_sequence(sentence, max_seq_length) for sentence in sentences]
# 构建词汇表，并为每个词汇赋予一个唯一的索引
# 这段代码将所有的文本合并成一个字符串，然后按照空格分割成词汇列表。
# 使用set去重，确保每个词汇只出现一次，然后为每个词汇赋予一个唯一的索引。
# 构建出词汇表word_dict，其中键是词汇，值是对应的唯一索引，vocab_size表示词汇表的大小
word_list = " ".join(sentences).split()
word_list = list(set(word_list))
word_dict = {w: i for i, w in enumerate(word_list)}
vocab_size = len(word_dict)
# 创建 BiLSTM_Attention 类的实例 model
# 创建了一个带有注意力机制的双向LSTM文本分类模型
model = BiLSTM_Attention()
# 定义损失函数 criterion 为交叉熵损失，优化器 optimizer 为Adam优化器
# 交叉熵损失适用于多分类任务，而Adam优化器是一种常用的优化算法
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 准备输入数据 inputs 和目标标签 targets
# 这段代码首先将每个句子转换为对应的索引序列，并将其转换为PyTorch的LongTensor类型，作为输入数据inputs。
# targets是标签序列，也转换为PyTorch的LongTensor类型。
inputs = torch.LongTensor([np.asarray([word_dict[n] for n in sen.split()]) for sen in sentences])
targets = torch.LongTensor([out for out in labels]) # To using Torch Softmax Loss function
# 进行模型训练
for epoch in range(5000):
# 在每一轮训练中，首先将优化器的梯度清零（optimizer.zero_grad()）
optimizer.zero_grad()
# 然后将输入数据输入模型，得到模型的输出和注意力权重。
output, attention = model(inputs)
# 计算交叉熵损失
loss = criterion(output, targets)
if (epoch + 1) % 1000 == 0:
print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.6f}'.format(loss))
# 进行反向传播
loss.backward()
# 参数优化
optimizer.step()
# Test
# 定义一个测试文本 test_text
test_text = '她 不高兴'
# 并转换为张量
tests = [np.asarray([word_dict[n] for n in test_text.split()])]
# print(tests)
test_batch = torch.LongTensor(tests)
# 进行预测
predict, _ = model(test_batch)
predict = predict.data.max(1, keepdim=True)[1]
# 根据预测结果，判断测试文本的意义是"好"还是"不好"
if predict[0][0] == 0:
print(test_text,"is Bad Mean...")
else:
print(test_text,"is Good Mean!!")
# 可视化注意力权重
# 这段代码使用Matplotlib绘制了注意力权重矩阵的热力图。
# 矩阵的横轴表示文本序列中的每个词汇（"first_word", "second_word", "third_word"）
# 纵轴表示输入数据的批次（"batch_1", "batch_2", "batch_3", "batch_4", "batch_5", "batch_6"）。
# 不同颜色的方块表示不同位置的词汇在分类时所受到的注意力程度。
# 通过该热力图可以观察模型在分类时关注的重要词
fig = plt.figure(figsize=(6, 3)) # [batch_size, n_step]
ax = fig.add_subplot(1, 1, 1)
ax.matshow(attention, cmap='viridis')
ax.set_xticklabels(['']+['first_word', 'second_word', 'third_word'], fontdict={'fontsize': 14}, rotation=90)
ax.set_yticklabels(['']+['batch_1', 'batch_2', 'batch_3', 'batch_4', 'batch_5', 'batch_6'], fontdict={'fontsize': 14})
plt.show()
