-
Notifications
You must be signed in to change notification settings - Fork 138
/
crfnp_dependency.py
198 lines (169 loc) · 8.22 KB
/
crfnp_dependency.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
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
from supar.models import CRFNPDependencyModel
from supar.parsers.biaffine_dependency import BiaffineDependencyParser
from supar.utils import Config
from supar.utils.logging import get_logger, progress_bar
from supar.utils.metric import AttachmentMetric
logger = get_logger(__name__)
class CRFNPDependencyParser(BiaffineDependencyParser):
r"""
The implementation of non-projective CRF Dependency Parser.
References:
- Xuezhe Ma and Eduard Hovy. 2017.
`Neural Probabilistic Model for Non-projective MST Parsing`_.
- Terry Koo, Amir Globerson, Xavier Carreras and Michael Collins. 2007.
`Structured Prediction Models via the Matrix-Tree Theorem`_.
.. _Neural Probabilistic Model for Non-projective MST Parsing:
https://www.aclweb.org/anthology/I17-1007/
.. _Structured Prediction Models via the Matrix-Tree Theorem:
https://www.aclweb.org/anthology/D07-1015/
"""
NAME = 'crfnp-dependency'
MODEL = CRFNPDependencyModel
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def train(self, train, dev, test, buckets=32, batch_size=5000, punct=False,
mbr=True, tree=False, proj=False, verbose=True, **kwargs):
r"""
Args:
train/dev/test (list[list] or str):
Filenames of the train/dev/test datasets.
buckets (int):
The number of buckets that sentences are assigned to. Default: 32.
batch_size (int):
The number of tokens in each batch. Default: 5000.
punct (bool):
If ``False``, ignores the punctuations during evaluation. Default: ``False``.
mbr (bool):
If ``True``, returns marginals for MBR decoding. Default: ``True``.
tree (bool):
If ``True``, ensures to output well-formed trees. Default: ``False``.
proj (bool):
If ``True``, ensures to output projective trees. Default: ``False``.
partial (bool):
``True`` denotes the trees are partially annotated. Default: ``False``.
verbose (bool):
If ``True``, increases the output verbosity. Default: ``True``.
kwargs (dict):
A dict holding the unconsumed arguments that can be used to update the configurations for training.
"""
return super().train(**Config().update(locals()))
def evaluate(self, data, buckets=8, batch_size=5000, punct=False,
mbr=True, tree=True, proj=False, verbose=True, **kwargs):
r"""
Args:
data (str):
The data for evaluation, both list of instances and filename are allowed.
buckets (int):
The number of buckets that sentences are assigned to. Default: 32.
batch_size (int):
The number of tokens in each batch. Default: 5000.
punct (bool):
If ``False``, ignores the punctuations during evaluation. Default: ``False``.
mbr (bool):
If ``True``, returns marginals for MBR decoding. Default: ``True``.
tree (bool):
If ``True``, ensures to output well-formed trees. Default: ``False``.
proj (bool):
If ``True``, ensures to output projective trees. Default: ``False``.
partial (bool):
``True`` denotes the trees are partially annotated. Default: ``False``.
verbose (bool):
If ``True``, increases the output verbosity. Default: ``True``.
kwargs (dict):
A dict holding the unconsumed arguments that can be used to update the configurations for evaluation.
Returns:
The loss scalar and evaluation results.
"""
return super().evaluate(**Config().update(locals()))
def predict(self, data, pred=None, buckets=8, batch_size=5000, prob=False,
mbr=True, tree=True, proj=False, verbose=True, **kwargs):
r"""
Args:
data (list[list] or str):
The data for prediction, both a list of instances and filename are allowed.
pred (str):
If specified, the predicted results will be saved to the file. Default: ``None``.
buckets (int):
The number of buckets that sentences are assigned to. Default: 32.
batch_size (int):
The number of tokens in each batch. Default: 5000.
prob (bool):
If ``True``, outputs the probabilities. Default: ``False``.
mbr (bool):
If ``True``, returns marginals for MBR decoding. Default: ``True``.
tree (bool):
If ``True``, ensures to output well-formed trees. Default: ``False``.
proj (bool):
If ``True``, ensures to output projective trees. Default: ``False``.
verbose (bool):
If ``True``, increases the output verbosity. Default: ``True``.
kwargs (dict):
A dict holding the unconsumed arguments that can be used to update the configurations for prediction.
Returns:
A :class:`~supar.utils.Dataset` object that stores the predicted results.
"""
return super().predict(**Config().update(locals()))
def _train(self, loader):
self.model.train()
bar, metric = progress_bar(loader), AttachmentMetric()
for words, feats, arcs, rels in bar:
self.optimizer.zero_grad()
mask = words.ne(self.WORD.pad_index)
# ignore the first token of each sentence
mask[:, 0] = 0
s_arc, s_rel = self.model(words, feats)
loss, s_arc = self.model.loss(s_arc, s_rel, arcs, rels, mask, self.args.mbr)
loss.backward()
nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip)
self.optimizer.step()
self.scheduler.step()
arc_preds, rel_preds = self.model.decode(s_arc, s_rel, mask)
# ignore all punctuation if not specified
if not self.args.punct:
mask &= words.unsqueeze(-1).ne(self.puncts).all(-1)
metric(arc_preds, rel_preds, arcs, rels, mask)
bar.set_postfix_str(f"lr: {self.scheduler.get_last_lr()[0]:.4e} - loss: {loss:.4f} - {metric}")
@torch.no_grad()
def _evaluate(self, loader):
self.model.eval()
total_loss, metric = 0, AttachmentMetric()
for words, feats, arcs, rels in loader:
mask = words.ne(self.WORD.pad_index)
# ignore the first token of each sentence
mask[:, 0] = 0
s_arc, s_rel = self.model(words, feats)
loss, s_arc = self.model.loss(s_arc, s_rel, arcs, rels, mask, self.args.mbr)
arc_preds, rel_preds = self.model.decode(s_arc, s_rel, mask)
# ignore all punctuation if not specified
if not self.args.punct:
mask &= words.unsqueeze(-1).ne(self.puncts).all(-1)
total_loss += loss.item()
metric(arc_preds, rel_preds, arcs, rels, mask)
total_loss /= len(loader)
return total_loss, metric
@torch.no_grad()
def _predict(self, loader):
self.model.eval()
preds = {}
arcs, rels, probs = [], [], []
for words, feats in progress_bar(loader):
mask = words.ne(self.WORD.pad_index)
# ignore the first token of each sentence
mask[:, 0] = 0
lens = mask.sum(1).tolist()
s_arc, s_rel = self.model(words, feats)
arc_preds, rel_preds = self.model.decode(s_arc, s_rel, mask)
arcs.extend(arc_preds[mask].split(lens))
rels.extend(rel_preds[mask].split(lens))
if self.args.prob:
arc_probs = s_arc if self.args.mbr else s_arc.softmax(-1)
probs.extend([prob[1:i+1, :i+1].cpu() for i, prob in zip(lens, arc_probs.unbind())])
arcs = [seq.tolist() for seq in arcs]
rels = [self.REL.vocab[seq.tolist()] for seq in rels]
preds = {'arcs': arcs, 'rels': rels}
if self.args.prob:
preds['probs'] = probs
return preds