-
Notifications
You must be signed in to change notification settings - Fork 0
/
C99.py
156 lines (143 loc) · 6.37 KB
/
C99.py
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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
# A python implementation of C99 algorithm for topic segmentation
# Modified from https://github.com/SALT-NLP/Multi-View-Seq2Seq/tree/master/data
import numpy as np
import numpy.typing as npt
from typing import Optional
class C99:
"""
C99 segmentation algorithm, see paper for more details
Reference:
"Advances in domain independent linear text segmentation"
"""
def __init__(self, window: int = 4, std_coeff: float = 1.2):
"""
window:window size for local similarity ranking
std_coeff: threshold to determine boundary, see paper for more details
"""
self.window = window
self.sim: Optional[npt.NDArray[np.float64]] = None
self.rank: Optional[npt.NDArray[np.float64]] = None
self.sm:Optional[npt.NDArray[np.float64]] = None
self.std_coeff = std_coeff
def segment(self, document_embeddings: npt.NDArray[np.float64]) -> npt.NDArray[np.int64]:
"""
document_embeddings: array with embeddings of document sentences (or other segmentation unit)
return list of zeros and ones, i-th element denotes whether exists a boundary right before paragraph i(0 indexed)
"""
# do not segment documents with less than 3 embeddings
if document_embeddings.shape[0] <= 3:
out = np.zeros(document_embeddings.shape[0], dtype=np.int64)
out[0] = 1
return out
# step 1, preprocessing
n = document_embeddings.shape[0]
self.window = min(self.window, n)
# step 2, compute similarity matrix
self.sim = np.zeros((n, n))
for i in range(n):
for j in range(i, n):
norm = np.linalg.norm(document_embeddings[i]) * np.linalg.norm(document_embeddings[j])
cosine_similarity = np.dot(document_embeddings[i], document_embeddings[j]) / norm
self.sim[i][j] = self.sim[j][i] = cosine_similarity
# step 3, compute rank matrix & sum matrix
self.rank = np.zeros((n, n))
for i in range(n):
for j in range(i, n):
r1 = max(0, i - self.window + 1)
r2 = min(n - 1, i + self.window - 1)
c1 = max(0, j - self.window + 1)
c2 = min(n - 1, j + self.window - 1)
sublist = self.sim[r1:(r2 + 1), c1:(c2+1)].flatten()
lowlist = [x for x in sublist if x < self.sim[i][j]]
self.rank[i][j] = 1.0 * len(lowlist) / ((r2 - r1 + 1) * (c2 - c1 + 1))
self.rank[j][i] = self.rank[i][j]
self.sm = np.zeros((n, n))
prefix_sm = np.zeros((n, n))
for i in range(n):
for j in range(n):
prefix_sm[i][j] = self.rank[i][j]
if i - 1 >= 0: prefix_sm[i][j] += prefix_sm[i - 1][j]
if j - 1 >= 0: prefix_sm[i][j] += prefix_sm[i][j - 1]
if i - 1 >= 0 and j - 1 >= 0: prefix_sm[i][j] -= prefix_sm[i - 1][j - 1]
for i in range(n):
for j in range(i, n):
if i == 0:
self.sm[i][j] = prefix_sm[j][j]
else:
self.sm[i][j] = prefix_sm[j][j] - prefix_sm[i - 1][j] \
- prefix_sm[j][i - 1] + prefix_sm[i - 1][i - 1]
self.sm[j][i] = self.sm[i][j]
# step 4, determine boundaries
D = 1.0 * self.sm[0][n - 1] / (n * n)
darr, region_arr, idx = [D], [_Region(0, n - 1, self.sm)], []
sum_region, sum_area = float(self.sm[0][n - 1]), float(n * n)
for i in range(n - 1):
mx, pos = -1e9, -1
for j, region in enumerate(region_arr):
if region.l == region.r:
continue
region.split(self.sm)
den = sum_area - region.area + region.lch.area + region.rch.area
cur = (sum_region - region.tot + region.lch.tot + region.rch.tot) / den
if cur > mx:
mx, pos = cur, j
assert(pos >= 0)
tmp = region_arr[pos]
region_arr[pos] = tmp.rch
region_arr.insert(pos, tmp.lch)
sum_region += tmp.lch.tot + tmp.rch.tot - tmp.tot
sum_area += tmp.lch.area + tmp.rch.area - tmp.area
darr.append(sum_region / sum_area)
idx.append(tmp.best_pos)
dgrad = [(darr[i + 1] - darr[i]) for i in range(len(darr) - 1)]
# optional step, smooth gradient
smooth_dgrad = [dgrad[i] for i in range(len(dgrad))]
if len(dgrad) > 1:
smooth_dgrad[0] = (dgrad[0] * 2 + dgrad[1]) / 3.0
smooth_dgrad[-1] = (dgrad[-1] * 2 + dgrad[-2]) / 3.0
for i in range(1, len(dgrad) - 1):
smooth_dgrad[i] = (dgrad[i - 1] + 2 * dgrad[i] + dgrad[i + 1]) / 4.0
dgrad = smooth_dgrad
avg, stdev = np.average(dgrad), np.std(dgrad)
cutoff = avg + self.std_coeff * stdev
assert(len(idx) == len(dgrad))
above_cutoff_idx = [i for i in range(len(dgrad)) if dgrad[i] >= cutoff]
if len(above_cutoff_idx) == 0: boundary = []
else: boundary = idx[:max(above_cutoff_idx) + 1]
ret = [0 for _ in range(n)]
for i in boundary:
ret[i] = 1
# boundary should not be too close
for j in range(i - 1, i + 2):
if j >= 0 and j < n and j != i and ret[j] == 1:
ret[i] = 0
break
return np.array([1] + ret[:-1], dtype=np.int64)
class _Region:
"""
Used to denote a rectangular region of similarity matrix.
"""
def __init__(self, l, r, sm_matrix):
assert(r >= l)
self.tot = sm_matrix[l][r]
self.l = l
self.r = r
self.area = (r - l + 1)**2
self.lch, self.rch, self.best_pos = None, None, -1
def split(self, sm_matrix):
if self.best_pos >= 0:
return
if self.l == self.r:
self.best_pos = self.l
return
assert(self.r > self.l)
mx, pos = -1e9, -1
for i in range(self.l, self.r):
carea = (i - self.l + 1)**2 + (self.r - i)**2
cur = (sm_matrix[self.l][i] + sm_matrix[i + 1][self.r]) / carea
if cur > mx:
mx, pos = cur, i
assert(pos >= self.l and pos < self.r)
self.lch = _Region(self.l, pos, sm_matrix)
self.rch = _Region(pos + 1, self.r, sm_matrix)
self.best_pos = pos