skip-thoughts/process_wiki.py at master · junfenglx/skip-thoughts · GitHub

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
# -*- coding: utf-8 -*-

# Created by junfeng on 3/28/16.

# logging config
import logging

logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s',
                    datefmt='%m/%d/%Y %I:%M:%S %p',
                    level=logging.INFO)
logger = logging.getLogger(__name__)

import os.path
import sys

from gensim.corpora import WikiCorpus

if __name__ == '__main__':

    logger.info("Running %s" % ' '.join(sys.argv))

    # check and process input arguments
    if len(sys.argv) < 3:
        print(sys.argv[0] + ' wiki-pages-articles.xml.bz2 output_text_name')
        sys.exit(1)
    inp, outp = sys.argv[1:3]
    space = ' '
    i = 0

    with open(outp, 'w') as output:
        wiki = WikiCorpus(inp, lemmatize=False, dictionary={})
        for text in wiki.get_texts():
            output.write(space.join(text) + '\n')
            i += 1
            if (i % 10000 == 0):
                logger.info('Saved ' + str(i) + ' articles')

        logger.info('Finished saved ' + str(i) + ' articles')