wiki.py 1.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940
  1. """
  2. Index Wikipedia
  3. """
  4. import gzip
  5. import html
  6. from urllib.parse import quote
  7. from spacy.lang.en import English
  8. from index import TinyIndexer, index_titles_urls_and_extracts, PAGE_SIZE, NUM_PAGES
  9. from paths import WIKI_TITLES_PATH, INDEX_PATH
  10. TEXT_TAGS = ['mediawiki', 'page', 'revision', 'text']
  11. TITLE_START = '<title>Wikipedia: '
  12. TITLE_END = '</title>\n'
  13. def index_wiki():
  14. nlp = English()
  15. with TinyIndexer(INDEX_PATH, NUM_PAGES, PAGE_SIZE) as indexer:
  16. titles_and_urls = get_wiki_titles_and_urls()
  17. index_titles_urls_and_extracts(indexer, nlp, titles_and_urls)
  18. def get_wiki_titles_and_urls():
  19. start_len = len(TITLE_START)
  20. end_len = len(TITLE_END)
  21. with gzip.open(WIKI_TITLES_PATH, 'rt') as wiki_titles_file:
  22. wiki_titles_file.readline()
  23. for raw_title in wiki_titles_file:
  24. assert raw_title.startswith(TITLE_START)
  25. assert raw_title.endswith(TITLE_END)
  26. title = raw_title[start_len:-end_len]
  27. unescaped_title = html.unescape(title)
  28. url = 'https://en.wikipedia.org/wiki/' + quote(unescaped_title.replace(' ', '_'))
  29. yield unescaped_title, url
  30. if __name__ == '__main__':
  31. index_wiki()