SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Added proper context management to ensure cleanup happens in the right

Browse Source 
order.
This commit is contained in:
Matthew Carrigan
2019-03-21 17:50:03 +00:00
parent 06a30cfdf3
commit abb7d1ff6d
1 changed files with 39 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
examples/lm_finetuning/pregenerate_training_data.py
View File

@@ -23,6 +23,7 @@ class DocumentDatabase:
self.documents = [] self.documents = []
self.document_shelf = None self.document_shelf = None
self.document_shelf_filepath = None self.document_shelf_filepath = None
self.temp_dir = None
self.doc_lengths = [] self.doc_lengths = []
self.doc_cumsum = None self.doc_cumsum = None
self.cumsum_max = None self.cumsum_max = None
@@ -68,9 +69,14 @@ class DocumentDatabase:
else: else:
return self.documents[item] return self.documents[item]
def cleanup(self): def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, traceback):
if self.document_shelf is not None: if self.document_shelf is not None:
self.document_shelf.close() self.document_shelf.close()
if self.temp_dir is not None:
self.temp_dir.cleanup()
def truncate_seq_pair(tokens_a, tokens_b, max_num_tokens): def truncate_seq_pair(tokens_a, tokens_b, max_num_tokens):
@@ -247,7 +253,7 @@ def main():
tokenizer = BertTokenizer.from_pretrained(args.bert_model, do_lower_case=args.do_lower_case) tokenizer = BertTokenizer.from_pretrained(args.bert_model, do_lower_case=args.do_lower_case)
vocab_list = list(tokenizer.vocab.keys()) vocab_list = list(tokenizer.vocab.keys())
docs = DocumentDatabase(reduce_memory=args.reduce_memory) with DocumentDatabase(reduce_memory=args.reduce_memory) as docs:
with args.train_corpus.open() as f: with args.train_corpus.open() as f:
doc = [] doc = []
for line in tqdm(f, desc="Loading Dataset", unit=" lines"): for line in tqdm(f, desc="Loading Dataset", unit=" lines"):
@@ -280,7 +286,6 @@ def main():
"max_seq_len": args.max_seq_len "max_seq_len": args.max_seq_len
} }
metrics_file.write(json.dumps(metrics)) metrics_file.write(json.dumps(metrics))
docs.cleanup()
if __name__ == '__main__': if __name__ == '__main__':