-
Notifications
You must be signed in to change notification settings - Fork 291
/
utils.py
executable file
·415 lines (334 loc) · 14.6 KB
/
utils.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
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
# -*- coding: utf-8 -*-
"""Useful functions for the pke module."""
from __future__ import division
from __future__ import absolute_import
from __future__ import print_function
import os
import sys
import csv
import pickle
import gzip
import json
import codecs
import logging
from collections import defaultdict
from pke.base import LoadFile
from pke.lang import stopwords, langcodes
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.decomposition import LatentDirichletAllocation
from nltk.stem.snowball import SnowballStemmer
def load_document_frequency_file(input_file,
delimiter='\t'):
"""Load a tsv (tab-separated-values) file containing document frequencies.
Automatically detects if input file is compressed (gzip) by looking at its
extension (.gz).
Args:
input_file (str): the input file containing document frequencies in
csv format.
delimiter (str): the delimiter used for separating term-document
frequencies tuples, defaults to '\t'.
Returns:
dict: a dictionary of the form {term_1: freq}, freq being an integer.
"""
# initialize the DF dictionary
frequencies = {}
# open the input file
with (gzip.open(input_file, 'rt', encoding='utf-8')
if input_file.endswith('.gz')
else codecs.open(input_file, 'rt', encoding='utf-8')) as f:
# read the csv file
df_reader = csv.reader(f, delimiter=delimiter)
# populate the dictionary
for row in df_reader:
frequencies[row[0]] = int(row[1])
# return the populated dictionary
return frequencies
def compute_document_frequency(documents,
output_file,
language='en',
stoplist=None,
normalization='stemming',
delimiter='\t',
# TODO: What is the use case for changing this ?
n=3):
"""Compute the n-gram document frequencies from a set of input documents.
An extra row is added to the output file for specifying the number of
documents from which the document frequencies were computed
(--NB_DOC-- tab XXX). The output file is compressed using gzip.
Args:
documents (list): list of pke-readable documents.
output_file (str): the output file.
language (str): language of the input documents (used for computing the
n-stem or n-lemma forms), defaults to 'en' (english).
stoplist (list): the stop words for filtering n-grams, default to
pke.lang.stopwords[language].
normalization (str): word normalization method, defaults to
'stemming'. Other possible value is 'none' for using word surface
forms instead of stems/lemmas.
delimiter (str): the delimiter between n-grams and document
frequencies, defaults to tabulation (\t).
n (int): the size of the n-grams, defaults to 3.
"""
# document frequency container
frequencies = defaultdict(int)
# initialize number of documents
nb_documents = 0
# loop through the documents
for document in documents:
# initialize load file object
doc = LoadFile()
# read the input file
doc.load_document(input=document,
language=language,
stoplist=stoplist,
normalization=normalization)
# candidate selection
doc.ngram_selection(n=n)
# filter candidates containing punctuation marks
doc.candidate_filtering()
# loop through candidates
for lexical_form in doc.candidates:
frequencies[lexical_form] += 1
nb_documents += 1
if nb_documents % 1000 == 0:
logging.info("{} docs, memory used: {} mb".format(
nb_documents,
sys.getsizeof(frequencies) / 1024 / 1024))
# create directories from path if not exists
if os.path.dirname(output_file):
os.makedirs(os.path.dirname(output_file), exist_ok=True)
# dump the df container
with gzip.open(output_file, 'wt', encoding='utf-8') as f:
# add the number of documents as special token
first_line = '--NB_DOC--' + delimiter + str(nb_documents)
f.write(first_line + '\n')
for ngram in frequencies:
line = ngram + delimiter + str(frequencies[ngram])
f.write(line + '\n')
def train_supervised_model(documents,
references,
model_file,
language='en',
stoplist=None,
normalization="stemming",
df=None,
model=None,
leave_one_out=False):
"""Build a supervised keyphrase extraction model from a set of documents
and reference keywords.
Args:
documents (list): list of tuple (id, pke-readable documents). `id`s
should match the one in reference.
references (dict): reference keywords.
model_file (str): the model output file.
language (str): language of the input documents (used for computing the
n-stem or n-lemma forms), defaults to 'en' (english).
stoplist (list): the stop words for filtering n-grams, default to
pke.lang.stopwords[language].
normalization (str): word normalization method, defaults to 'stemming'.
Other possible values are 'lemmatization' or 'None' for using word
surface forms instead of stems/lemmas.
df (dict): df weights dictionary.
model (object): the supervised model to train, defaults to Kea.
leave_one_out (bool): whether to use a leave-one-out procedure for
training, creating one model per input, defaults to False.
"""
training_instances = []
training_classes = []
masks = {}
# get the input files from the input directory
for doc_id, document in documents:
# logging.info('reading file {}'.format(input_file))
# get the document id from file name
# doc_id = '.'.join(os.path.basename(input_file).split('.')[0:-1])
# initialize the input file
model.__init__()
# load the document
model.load_document(input=document,
language=language,
stoplist=stoplist,
normalization=normalization)
# candidate selection
model.candidate_selection()
# skipping documents without candidates
if not len(model.candidates):
continue
# extract features
model.feature_extraction(df=df, training=True)
# add the first offset for leave-one-out masking
masks[doc_id] = [len(training_classes)]
# annotate the reference keyphrases in the instances
for candidate in model.instances:
if candidate in references[doc_id]:
training_classes.append(1)
else:
training_classes.append(0)
training_instances.append(model.instances[candidate])
# add the last offset for leave-one-out masking
masks[doc_id].append(len(training_classes))
if not leave_one_out:
logging.info('writing model to {}'.format(model_file))
model.train(training_instances=training_instances,
training_classes=training_classes,
model_file=model_file)
else:
logging.info('leave-one-out training procedure')
for doc_id in masks:
logging.info('writing model to {}'.format(doc_id))
ind = masks[doc_id]
fold = training_instances[:ind[0]] + training_instances[ind[1]:]
gold = training_classes[:ind[0]] + training_classes[ind[1]:]
model.train(training_instances=fold,
training_classes=gold,
model_file='{}.{}.pickle'.format(model_file, doc_id))
def load_references(input_file,
sep_doc_id=':',
sep_ref_keyphrases=',',
normalize_reference=False,
language="en",
encoding=None,
excluded_file=None):
"""Load a reference file. Reference file can be either in json format or in
the SemEval-2010 official format.
Args:
input_file (str): path to the reference file.
sep_doc_id (str): the separator used for doc_id in reference file,
defaults to ':'.
sep_ref_keyphrases (str): the separator used for keyphrases in
reference file, defaults to ','.
normalize_reference (bool): whether to normalize the reference
keyphrases using stemming, default to False.
language (str): language of the input documents (used for computing the
stems), defaults to 'en' (english).
encoding (str): file encoding, default to None.
excluded_file (str): file to exclude (for leave-one-out
cross-validation), defaults to None.
"""
logging.info('loading reference keyphrases from {}'.format(input_file))
references = defaultdict(list)
# open input file
with codecs.open(input_file, 'r', encoding) as f:
# load json data
if input_file.endswith('.json'):
references = json.load(f)
for doc_id in references:
references[doc_id] = [keyphrase for variants in
references[doc_id] for keyphrase in
variants]
# or load SemEval-2010 file
else:
for line in f:
cols = line.strip().split(sep_doc_id)
doc_id = cols[0].strip()
keyphrases = cols[1].strip().split(sep_ref_keyphrases)
for v in keyphrases:
if '+' in v:
for s in v.split('+'):
references[doc_id].append(s)
else:
references[doc_id].append(v)
# normalize reference if needed
if normalize_reference:
# initialize stemmer
langcode = langcodes.get(language.replace('en', 'xx'), 'porter')
stemmer = SnowballStemmer(langcode)
for doc_id in references:
for i, keyphrase in enumerate(references[doc_id]):
stems = [stemmer.stem(w) for w in keyphrase.split()]
references[doc_id][i] = ' '.join(stems)
# remove excluded file if needed
if excluded_file is not None:
if excluded_file not in references:
logging.warning("{} is not in references".format(excluded_file))
else:
logging.info("{} removed from references".format(excluded_file))
del references[excluded_file]
return references
def load_lda_model(input_file):
"""Load a gzip file containing lda model.
Args:
input_file (str): the gzip input file containing lda model.
Returns:
dictionary: a dictionary of the form {term_1: freq}, freq being an
integer.
model: an initialized sklearn.decomposition.LatentDirichletAllocation
model.
"""
model = LatentDirichletAllocation()
with gzip.open(input_file, 'rb') as f:
(dictionary,
model.components_,
model.exp_dirichlet_component_,
model.doc_topic_prior_) = pickle.load(f)
return dictionary, model
def compute_lda_model(documents,
output_file,
n_topics=500,
language="en",
stoplist=None,
normalization="stemming"):
"""Compute a LDA model from a collection of documents. Latent Dirichlet
Allocation is computed using sklearn module.
Args:
documents (str): list fo pke-readable documents.
output_file (str): the output file.
n_topics (int): number of topics for the LDA model, defaults to 500.
language (str): language of the input documents, used for stop_words
in sklearn CountVectorizer, defaults to 'en'.
stoplist (list): the stop words for filtering words, default to
pke.lang.stopwords[language].
normalization (str): word normalization method, defaults to
'stemming'. Other possible value is 'none'
for using word surface forms instead of stems/lemmas.
"""
# texts container
texts = []
# loop throught the documents
for document in documents:
# initialize load file object
doc = LoadFile()
# read the input file
doc.load_document(input=document,
language=language,
normalization=normalization)
# container for current document
text = []
# loop through sentences
for sentence in doc.sentences:
# get the tokens (stems) from the sentence if they are not
# punctuation marks
text.extend([sentence.stems[i] for i in range(sentence.length)
if sentence.pos[i] != 'PUNCT'
and sentence.pos[i].isalpha()])
# add the document to the texts container
texts.append(' '.join(text))
# vectorize dataset
# get the stoplist from pke.lang because CountVectorizer only contains
# english stopwords atm
if stoplist is None:
# CountVectorizer expects a list
# stopwords.get is a set
stoplist = list(stopwords.get(language))
tf_vectorizer = CountVectorizer(
stop_words=stoplist)
tf = tf_vectorizer.fit_transform(texts)
# extract vocabulary
vocabulary = list(tf_vectorizer.get_feature_names_out())
# create LDA model and train
lda_model = LatentDirichletAllocation(n_components=n_topics,
random_state=0,
learning_method='batch')
lda_model.fit(tf)
# save all data necessary for later prediction
saved_model = (vocabulary,
lda_model.components_,
lda_model.exp_dirichlet_component_,
lda_model.doc_topic_prior_)
# Dump the df container
logging.info('writing LDA model to {}'.format(output_file))
# create directories from path if not exists
if os.path.dirname(output_file):
os.makedirs(os.path.dirname(output_file), exist_ok=True)
# dump the LDA model
with gzip.open(output_file, 'wb') as fp:
pickle.dump(saved_model, fp)