wikipedia_preprocessing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import simplejson as json\n",
    "import time\n",
    "import smart_open\n",
    "import re\n",
    "import sys\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!python -m gensim.scripts.segment_wiki -i -f enwiki-20200101-pages-articles-multistream1.xml-p10p30302.bz2 -o enwiki-latest.json.gz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "RE_P0 = re.compile('<!--.*?-->', re.DOTALL | re.UNICODE)  # comments\n",
    "RE_P1 = re.compile('<ref([> ].*?)(</ref>|/>)', re.DOTALL | re.UNICODE)  # footnotes\n",
    "RE_P2 = re.compile(\"(\\n\\[\\[[a-z][a-z][\\w-]*:[^:\\]]+\\]\\])+$\", re.UNICODE)  # links to languages\n",
    "RE_P3 = re.compile(\"{{([^}{]*)}}\", re.DOTALL | re.UNICODE)  # template\n",
    "RE_P4 = re.compile(\"{{([^}]*)}}\", re.DOTALL | re.UNICODE)  # template\n",
    "RE_P5 = re.compile('\\[(\\w+):\\/\\/(.*?)(( (.*?))|())\\]', re.UNICODE)  # remove URL, keep description\n",
    "RE_P6 = re.compile(\"\\[([^][]*)\\|([^][]*)\\]\", re.DOTALL | re.UNICODE)  # simplify links, keep description\n",
    "RE_P7 = re.compile('\\n\\[\\[[iI]mage(.*?)(\\|.*?)*\\|(.*?)\\]\\]', re.UNICODE)  # keep description of images\n",
    "RE_P8 = re.compile('\\n\\[\\[[fF]ile(.*?)(\\|.*?)*\\|(.*?)\\]\\]', re.UNICODE)  # keep description of files\n",
    "RE_P9 = re.compile('<nowiki([> ].*?)(</nowiki>|/>)', re.DOTALL | re.UNICODE)  # outside links\n",
    "RE_P10 = re.compile('<math([> ].*?)(</math>|/>)', re.DOTALL | re.UNICODE)  # math content\n",
    "RE_P11 = re.compile('<(.*?)>', re.DOTALL | re.UNICODE)  # all other tags\n",
    "RE_P12 = re.compile('\\n(({\\|)|(\\|-)|(\\|}))(.*?)(?=\\n)', re.UNICODE)  # table formatting\n",
    "RE_P13 = re.compile('\\n(\\||\\!)(.*?\\|)*([^|]*?)', re.UNICODE)  # table cell formatting\n",
    "RE_P14 = re.compile('\\[\\[Category:[^][]*\\]\\]', re.UNICODE)  # categories\n",
    "# Remove File and Image template\n",
    "RE_P15 = re.compile('\\[\\[([fF]ile:|[iI]mage)[^]]*(\\]\\])', re.UNICODE)\n",
    "RE_P16 = re.compile('=*\\**', re.UNICODE)\n",
    "\n",
    "def remove_markup(text):\n",
    "    text = re.sub(RE_P2, \"\", text)  # remove the last list (=languages)\n",
    "    iters = 0\n",
    "    while True:\n",
    "        old, iters = text, iters + 1\n",
    "        text = re.sub(RE_P0, \"\", text)  # remove comments\n",
    "        text = re.sub(RE_P1, '', text)  # remove footnotes\n",
    "        text = re.sub(RE_P9, \"\", text)  # remove outside links\n",
    "        text = re.sub(RE_P10, \"\", text)  # remove math content\n",
    "        text = re.sub(RE_P11, \"\", text)  # remove all remaining tags\n",
    "        text = re.sub(RE_P14, '', text)  # remove categories\n",
    "        text = re.sub(RE_P5, '\\\\3', text)  # remove urls, keep description\n",
    "        text = re.sub(RE_P6, '\\\\2', text)  # simplify links, keep description only\n",
    "        # remove table markup\n",
    "        text = text.replace('||', '\\n|')  # each table cell on a separate line\n",
    "        text = re.sub(RE_P12, '\\n', text)  # remove formatting lines\n",
    "        text = re.sub(RE_P13, '\\n\\\\3', text)  # leave only cell content\n",
    "        text = re.sub(RE_P16, '', text)  # leave only cell content\n",
    "        # remove empty mark-up\n",
    "        text = text.replace('[]', '')\n",
    "        if old == text or iters > 2:  # stop if nothing changed between two iterations or after a fixed number of iterations\n",
    "            break\n",
    "    # the following is needed to make the tokenizer see '[[socialist]]s' as a single word 'socialists'\n",
    "    text = text.replace('[', '').replace(']', '')  # promote all remaining markup to plain text\n",
    "    return text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_article(title, text,  directory=\"./wikipedia/\"):\n",
    "    if not os.path.exists(directory):\n",
    "        try:\n",
    "            os.makedirs(directory)\n",
    "        except OSError as e:\n",
    "            print(e)\n",
    "        path=os.path.join(directory, title.replace(\" \",\"_\").replace(\"/\",\"_\") +'.txt')\n",
    "        if os.path.exists(path):\n",
    "            print('article already present')\n",
    "        else:\n",
    "            with open(path, 'w') as fp:\n",
    "                fp.write(text)\n",
    "    else:\n",
    "        path=os.path.join(directory, title.replace(\" \",\"_\").replace(\"/\",\"_\") +'.txt')\n",
    "        if os.path.exists(path):\n",
    "            print('article already present')\n",
    "        else:\n",
    "            with open(path, 'w') as fp:\n",
    "                fp.write(text)\n",
    "    return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__== \"__main__\":\n",
    "    filename=\"./enwiki-latest.json.gz\"\n",
    "    start_time=time.time()\n",
    "    count=0\n",
    "    try:\n",
    "        for line in smart_open.open(filename):\n",
    "            count+=1\n",
    "            article = json.loads(line)\n",
    "            title = article['title']\n",
    "            print(title)\n",
    "            text = remove_markup(\" \".join(article['section_texts']))\n",
    "            save_article(title, text)\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        print(\"In Exception\")\n",
    "        pass\n",
    "    print(\"Total articles: %d\" %count)\n",
    "    print(\"total time(in min): %f\"  %((time.time()-start_time)/60))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}