TSA / outline.py
QINGCHE's picture
fix bugs
8ba144e
raw
history blame
6.54 kB
import numpy as np
from scipy.cluster.hierarchy import linkage, fcluster, dendrogram
import matplotlib.pyplot as plt
import numpy as np
def find_parent(matrix, node):
# 获取该节点的所有可能的父节点及其概率
parents = matrix[:, node]
# 找出其中概率最大的父节点
max_parent = np.argmax(parents)
# 如果该父节点的概率不为零,那么返回该父节点
if parents[max_parent] > 0:
return max_parent
# 否则,返回None,表示该节点没有父节点
else:
return None
def find_tree(matrix, node, depth=1, children=[], max_depth=1, visited=set()):
# 定义一个空列表,用来存储结果
result = []
# 调用find_parent函数,找出该节点的最可能的父节点
parent = find_parent(matrix, node)
# 如果该父节点存在,且不在visited中
if parent is not None and parent not in visited:
# 将该父子关系添加到结果列表中
result.append([parent, node])
# 遍历当前节点的所有子节点
for child in children:
# 将该子节点和当前节点的关系也添加到结果列表中
result.append([node, child])
# 如果当前深度小于最大深度,那么递归地调用find_tree函数,传入概率矩阵和该父节点作为参数,并将深度加一
if depth < max_depth:
# 将当前节点添加到visited中
visited.add(node)
# 递归调用find_tree函数,并将visited传入
result.extend(find_tree(matrix, parent, depth + 1, visited=visited))
# 返回结果列表
return result
# 定义一个函数,它接受一个树形结构作为参数,然后返回该树形结构的概率
def find_prob(tree, matrix):
# 定义一个变量,用来存储概率,初始值为1
prob = 1
# 遍历树形结构中的每个父子关系
for parent, child in tree:
# 将该父子关系对应的概率乘到变量上
prob *= matrix[parent][child]
# 返回最终的概率
return prob
# 定义一个函数,它接受一个概率矩阵和一个k值作为参数,然后返回前k个最可能的森林及其概率
def find_forests(matrix, k):
# 定义一个空字典,用来存储每个森林及其概率
forests = {}
# 遍历所有的节点
for i in range(len(matrix)):
# 获取该节点的所有可能的子节点及其概率
children = matrix[i]
# 定义一个空列表,用来存储该节点的所有子节点
child_list = []
# 遍历所有的子节点
for j in range(len(children)):
# 如果子节点的概率不为零
if children[j] > 0:
# 将该子节点添加到列表中
child_list.append(j)
# 调用find_tree函数,传入概率矩阵、当前节点和子节点列表作为参数,并将返回的结果赋值给tree
tree = find_tree(matrix, i, children=child_list)
tree = tuple([tuple(x) for x in tree]) # 将列表中的每个元素转换为元组,然后再将整个列表转换为元组
if tree:
# 调用find_prob函数,传入tree作为参数,并将返回的结果赋值给prob
prob = find_prob(tree, matrix)
# 如果tree已经在字典的键中,即该树形结构已经被添加过
if tuple(tree) in forests: # 将tree转换为元组
# 那么将prob加到字典的值上,表示该树形结构出现了多次
forests[tuple(tree)] += prob # 将tree转换为元组
# 否则,即该树形结构是第一次被添加
else:
# 那么将tree和prob作为键值对添加到字典中
forests[tuple(tree)] = prob # 将tree转换为元组
# 将字典按照值降序排序,得到一个列表,每个元素是一个键值对的元组
sorted_forests = sorted(forests.items(), key=lambda x: x[1], reverse=True)
# 返回列表的第一个元素,即最可能的森林及其概率
forest, prob = sorted_forests[0]
# 定义一个空字典,用来存储带有层级结构的森林
result = {}
# 遍历森林中的每个树形结构
for parent, child in forest:
# 如果父节点已经在字典的键中,即该父节点已经有其他子节点
if parent in result:
# 那么将当前子节点添加到父节点对应的值中,即该父节点的子节点列表中
result[parent].append(child)
# 否则,即该父节点是第一次出现
else:
# 那么将父节点和当前子节点作为键值对添加到字典中,即创建一个新的子节点列表
result[parent] = [child]
# 返回带有层级结构的字典和概率
return result, prob
def passage_outline(matrix,sentences):
# Z = linkage(matrix, method="average")
# median = np.median(matrix)
# print(median)
# print('yes')
# labels = fcluster(Z, t=median, criterion="distance")
# print(labels)
# 调用函数,传入概率矩阵作为参数,并且指定最大深度为1
result, prob = find_forests(matrix, 1)
# 打印结果字典和概率
print(result, prob)
# 根据簇标签和主题句子生成文章结构
structure = {}
for each in result.keys():
structure[each] =[sentences[i] for i in result[each]]
# for label, sentence in zip(result.keys(), sentences):
# if label not in structure:
# structure[label] = []
# structure[label].append(sentence)
outline = ""
outline_list = []
for key in sorted(structure.keys()):
outline_list.append(f"主题:")
outline = outline+f"主题:\n"
for sentence in structure[key]:
outline_list.append(sentence)
outline = outline+f"- {sentence}\n"
return outline,outline_list
if __name__ == "__main__":
matrix = np.array([[0.0 ,0.02124888, 0.10647043 ,0.09494194 ,0.0689209 ],
[0.01600688 ,0.0 ,0.05879448 ,0.0331325 , 0.0155093 ],
[0.01491911 ,0.01652437, 0.0, 0.04714563, 0.04577385],
[0.01699071 ,0.0313585 , 0.040299 ,0.0 ,0.014933 ],
[0.02308992 ,0.02791895 ,0.06547201, 0.08517842 ,0.0]])
sentences = ["主题句子1", "主题句子2", "主题句子3", "主题句子4", "主题句子5"]
print(passage_outline(matrix,sentences)[0])