wiki_indexer.py

# -*- coding: utf-8 -*-
"""wikiPedia.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1NJzAYW0tLvN56LAuCempaOcZxBOUTRE6
"""


import sys
import time
import os
import xml.sax
import re
from collections import defaultdict
import math
import bisect
import nltk
from nltk.stem import *
from nltk.stem.snowball import SnowballStemmer
from nltk.tokenize import RegexpTokenizer
import pickle
from nltk.stem.porter import *
from nltk.stem import PorterStemmer as porter
from Stemmer import Stemmer
import operator

stopwords=["a", "about", "above", "above", "across", "after", "afterwards", "again",
           "against", "all", "almost", "alone", "along", "already", "also","although",
           "always","am","among", "amongst", "amoungst", "amount",  "an", "and", "another",
           "any","anyhow","anyone","anything","anyway", "anywhere", "are", "around", "as",
           "at", "back","be","became", "because","become","becomes", "becoming", "been",
           "before", "beforehand", "behind", "being", "below", "beside", "besides",
           "between", "beyond", "bill", "both", "bottom","but", "by", "call", "can",
           "cannot", "cant", "co", "con", "could", "couldnt", "cry", "de", "describe",
           "detail", "do", "done", "down", "due", "during", "each", "eg", "eight", "either",
           "eleven","else", "elsewhere", "empty", "enough", "etc", "even", "ever", "every",
           "everyone", "everything", "everywhere", "except", "few", "fifteen", "fify", "fill",
           "find", "fire", "first", "five", "for", "former", "formerly", "forty", "found",
           "four", "from", "front", "full", "further", "get", "give", "go", "had", "has",
           "hasnt", "have", "he", "hence", "her", "here", "hereafter", "hereby", "herein",
           "hereupon", "hers", "herself", "him", "himself", "his", "how", "however", "hundred",
           "ie", "if", "in", "inc", "indeed", "interest", "into", "is", "it", "its", "itself",
           "keep", "last", "latter", "latterly", "least", "less", "ltd", "made", "many", "may",
           "me", "meanwhile", "might", "mill", "mine", "more", "moreover", "most", "mostly",
           "move", "much", "must", "my", "myself", "name", "namely", "neither", "never",
           "nevertheless", "next", "nine", "no", "nobody", "none", "noone", "nor", "not",
           "nothing", "now", "nowhere", "of", "off", "often", "on", "once", "one", "only",
           "onto", "or", "other", "others", "otherwise", "our", "ours", "ourselves", "out",
           "over", "own","part", "per", "perhaps", "please", "put", "rather", "re", "same",
           "see", "seem", "seemed", "seeming", "seems", "serious", "several", "she", "should",
           "show", "side", "since", "sincere", "six", "sixty", "so", "some", "somehow", "someone",
           "something", "sometime", "sometimes", "somewhere", "still", "such", "system", "take",
           "ten", "than", "that", "the", "their", "them", "themselves", "then", "thence", "there",
           "thereafter", "thereby", "therefore", "therein", "thereupon", "these", "they", "thickv",
           "thin", "third", "this", "those", "though", "three", "through", "throughout", "thru",
           "thus", "to", "together", "too", "top", "toward", "towards", "twelve", "twenty", "two",
           "un", "under", "until", "up", "upon", "us", "very", "via", "was", "we", "well", "were",
           "what", "whatever", "when", "whence", "whenever", "where", "whereafter", "whereas",
           "whereby", "wherein", "whereupon", "wherever", "whether", "which", "while", "whither",
           "who", "whoever", "whole", "whom", "whose", "why", "will", "with", "within", "without",
           "would", "yet", "you", "your", "yours", "yourself", "yourselves"
]


# global variables
stemmer = Stemmer('porter')
indexTable  = {}
wordCount = 0
splitThreshold = 10000
wordInEachFileThreshold = 10000
initialFileCount = 0
folderPath = "index/"
docNo = 0
os.mkdir('titles')
titlePath = 'titles/' +  'titleMap.txt'
titlePickle = 'titles/' +  'titlePickle'
titleMap = {}
titleptr = open(titlePath,'w')





titlefileIdx = 0
titleThreshold = 100000
titleList = []



def writeTitleToFile():
  global titleList
  global titlefileIdx
  global titlePath
  if(len(titleList)):
    filePath = titlePath + str(titlefileIdx) + ".txt"
    fp1 = open(filePath,'w')
    for title in  titleList:
      fp1.write(title)
    fp1.close()
    titleList.clear()
    titlefileIdx += 1



def writeTitle(docTitle):
  global titleThreshold
  global titlePath
  titleList.append(docTitle)
  if(len(titleList) == titleThreshold):
    writeTitleToFile()
    


def body(data):
  global wordCount
  listOfBodyWords = cleanData(data)
  return listOfBodyWords

def tokenizeWords(data):
    tokenizer = RegexpTokenizer(r'[A-Za-z0-9]+')
    return tokenizer.tokenize(data)

def cleanData(data):
  global wordCount
  # Tokenisation -> lower -> stopWords -> stemming
  data = data.lower()
  data = tokenizeWords(data)
 
  words = []
  for token in data:
      wordCount += 1
      token = (token)
      if len(token) <= 1 or token in stopwords:
          continue
      words.append(stemmer.stemWord(token))
  return words



def infoBox(data):
  infoSection = data.split("{{Infobox")
  litsOfWordinInfobox = []
  if(len(infoSection) > 1):
    for line in range(1,len(infoSection)):
      eachWord = infoSection[line].split('\n')
      for word in eachWord:
        if( word == "}}"):
          break
        preProcessWord = cleanData(word)
        litsOfWordinInfobox += preProcessWord
  return litsOfWordinInfobox

def category(data):
   listOfWordsinCategory = []
   categorySection = re.findall(r"\[\[Category:(.*)\]\]", str(data))
   for word in categorySection:
     listOfWordsinCategory += cleanData(word)
   return listOfWordsinCategory

def external(data):
  externalLinksWords = []
  externalLinkSection = data.split("==External links==")
  if(len(externalLinkSection)) > 1:
    externalLinkSection = externalLinkSection[1].split("\n")
    for word in externalLinkSection:
      if  word:
        if word[0]=='*' :
          externalLinksWords += cleanData(word) 
  return externalLinksWords 

def reference(data):
  referenceListWords = []
  referenceLinkSection = data.split("==References==")
  if len(referenceLinkSection) > 1:
    referenceLinkSection = referenceLinkSection[1]
    referenceLinkSection = referenceLinkSection.split("\n")

    for references in referenceLinkSection:
      if ("[[Category" in references):
          break
      if ("==" in references):
        break
      if ("DEFAULTSORT" in references):
        break
        
      reference = cleanData(references)
      if "reflist" in reference:
          reference.remove("reflist")
      referenceListWords += reference
  
  return referenceListWords


def getAllsection(title,data):

  
  
  # get list of word belong to each section
  
  litsOfWordinInfobox = infoBox(data)
  listOfWordsinCategory  = category(data)
  externalLinksWords = external(data)
  referenceListWords = reference(data)
  listOfBodyWords   =   body(data)
  listOfTitlsWords = []

  for words in title:
    listOfTitlsWords += cleanData(title)


  return listOfTitlsWords , litsOfWordinInfobox , listOfWordsinCategory , externalLinksWords , referenceListWords , listOfBodyWords

def getMap(words):
  wordMap = {}
  for word in words:
    if word not in wordMap:
      wordMap[word] = 1
    else:
      wordMap[word] += 1
  return wordMap


def checkThreshold():
 
  global wordInEachFileThreshold
  global initialFileCount
  global folderPath
  if(len(indexTable)== (wordInEachFileThreshold+1)):
    # print("folderPath:",folderPath)
    index = "initialIndex_" + str(initialFileCount) + ".txt"
    # print(index)
    fileName = folderPath + index
    fp = open(fileName, "a")
    saveToText(fp)
    fp.close()
    initialFileCount += 1
    indexTable.clear()


def saveAndPreprocess(title , text , docId ):
  title , infoMap , categoryMap , externalData , referenceMap  , bodyMap= getAllsection(title , text)
  
  titleMap = getMap(title)
  categoryMap = getMap(categoryMap)
  bodyMap = getMap(bodyMap)
  infoMap = getMap(infoMap)
  referenceMap = getMap(referenceMap)
  externalData = getMap(externalData)

  
  for word in externalData:
    if word not in indexTable:
      indexTable[word] ={}
    if docId not in indexTable[word]:
      indexTable[word][docId] =  {}
    indexTable[word][docId]["externalData"] = externalData[word]
    checkThreshold()


  for word in referenceMap:
    if word not in indexTable:
      indexTable[word] ={}
    if docId not in indexTable[word]:
      indexTable[word][docId] = {}
    indexTable[word][docId]["referenceMap"] = referenceMap[word]
    checkThreshold()

  for word in infoMap:
    if word not in indexTable:
      indexTable[word] = {}
    if docId not in indexTable[word]:
      indexTable[word][docId] = {}
    indexTable[word][docId]["infoMap"] = infoMap[word]
    checkThreshold()

  for word in bodyMap:
    if word not in indexTable:
      indexTable[word] = {}
    if docId not in indexTable[word]:
      indexTable[word][docId] = {}
    indexTable[word][docId]["bodyMap"] = bodyMap[word]
    checkThreshold()


  for word in categoryMap:
    if word not in indexTable:
      indexTable[word] = {}
    if docId not in indexTable[word]:
      indexTable[word][docId] = {}
    indexTable[word][docId]["categoryMap"] = categoryMap[word]
    checkThreshold()


  for word in titleMap:
    if word not in indexTable:
      indexTable[word] = {}
    if docId not in indexTable[word]:
      indexTable[word][docId] = {}
    indexTable[word][docId]["titleMap"] = titleMap[word]
    checkThreshold()

def saveToText(fp):
  toText = None
  for word in sorted(indexTable):
    toText = word + ':'
    for docId in sorted(indexTable[word]):
      toText += 'd' + str(docId)
      
      if "infoMap" in indexTable[word][docId] and indexTable[word][docId]["infoMap"]>0:
        toText += 'i'
        toText += str(indexTable[word][docId]["infoMap"])
      
      if "referenceMap" in indexTable[word][docId] and indexTable[word][docId]["referenceMap"]>0:
        toText += 'r'
        toText += str(indexTable[word][docId]["referenceMap"])
      
      if "bodyMap" in indexTable[word][docId] and indexTable[word][docId]["bodyMap"]>0:
        toText += 'b'
        toText += str(indexTable[word][docId]["bodyMap"])
      
      if "categoryMap" in indexTable[word][docId] and indexTable[word][docId]["categoryMap"]>0:
        toText += 'c'
        toText += str(indexTable[word][docId]["categoryMap"])
      
      if "titleMap" in indexTable[word][docId] and indexTable[word][docId]["titleMap"]>0:
        toText += 't'
        toText += str(indexTable[word][docId]["titleMap"])

      toText += ' '
    
    fp.write(toText + "\n")



# def mergeFiles(folderLocationOfIndex):




class Documenthandler( xml.sax.ContentHandler ):
  global titleptr
  global docNo
  global titleMap
  def __init__(self):
   self.current_tag = ""
   self.title = ""
   self.docId = 0
   self.body = ""
  
   
  
  def startElement(self, tag, attributes):
   self.current_tag = tag
  
  def endElement(self, tag):
   if tag == "page":

     saveAndPreprocess(self.title,self.body , self.docId)
     titleMap[self.docId] = self.title.strip() 
     titleptr.write(self.title.strip())
     titleptr.write('\n')
     self.title = ""
     self.docId += 1
     docNo = self.docId
     self.body = ""
    
  
  def characters(self, content):
    if self.current_tag == "title":
      self.title += content
    if self.current_tag == "text":
      self.body += content


def split(absPathOfMergeFile,folderLocationOfIndex):
    global splitThreshold
  
    index = 0
    count = 0
    file1 = absPathOfMergeFile
    secIndFile = folderLocationOfIndex + "secondaryIndex.txt"
    
    fp1=open(file1,'r') 
    sf =open(secIndFile,'w')
    
    lineFile1 = fp1.readline().strip('\n')
    
    while( lineFile1 ):


        if(count == 0):
            indexFileName = folderLocationOfIndex + "index_" + str(index) + ".txt"
            wordFile1 = lineFile1.split(":")[0]
            sf.write(wordFile1 + '\n')
            fp2 = open(indexFileName, 'w')
        
        indexFileName = folderLocationOfIndex + "index_" + str(index) + ".txt"
        wordFile1 = lineFile1.split(":")[0]
        
        fp2.write(lineFile1 + '\n')
        count += 1
        
        if(count == splitThreshold):
            count = 0
            index += 1
            fp2.close()
        lineFile1 = fp1.readline().strip('\n')
    
    fp1.close()
    fp2.close()
    sf.close()
    os.remove(absPathOfMergeFile)



def mergeTwoFiles(file1 , file2 , folderLocationOfIndex):
    # print("merge:",file1," ",file2)
    if file1 == file2:
        return
    fp1 = open(file1, 'r')
    fp2 = open(file2, 'r')
    tempFile = folderLocationOfIndex + "temporary.txt" 
    fp3 = open(tempFile, 'w')
    lineFile1 = fp1.readline().strip('\n')
    lineFile2 = fp2.readline().strip('\n')
    while (lineFile1 and lineFile2):
        wordFile1 = lineFile1.split(":")[0]
        wordFile2 = lineFile2.split(":")[0]
        if wordFile2 < wordFile1:
            fp3.write(lineFile2 + '\n')
            lineFile2 = fp2.readline().strip('\n')
        elif wordFile1 < wordFile2:
            fp3.write(lineFile1 + '\n')
            lineFile1 = fp1.readline().strip('\n')
        else:
            list1 = lineFile1.strip().split(":")[1]
            list2 = lineFile2.strip().split(':')[1]
            fp3.write(wordFile1 + ':' + list1 + list2 + '\n')
            lineFile1 = fp1.readline().strip('\n')
            lineFile2 = fp2.readline().strip('\n')
    while lineFile1:
        fp3.write(lineFile1 + '\n')
        lineFile1 = fp1.readline().strip('\n')
    while lineFile2:
        fp3.write(lineFile2 + '\n')
        lineFile2 = fp2.readline().strip('\n')
    fp1.close()
    fp2.close()
    fp3.close()
    os.remove(file1)
    os.remove(file2)
    os.rename(tempFile, file1)



      
		

def mergeFiles(folderLocationOfIndex):
  if os.path.isdir(folderLocationOfIndex):
    for filename in os.listdir(folderLocationOfIndex):
      listOfIndexFiles = sorted(os.listdir(folderLocationOfIndex))
      while(len(listOfIndexFiles)>=2):
        for i in range(0,len(listOfIndexFiles),2):
          if (i+1 ) >= len(listOfIndexFiles):
            break
          file1 = folderLocationOfIndex + listOfIndexFiles[i]
          file2 = folderLocationOfIndex + listOfIndexFiles[i+1]
          mergeTwoFiles(file1 , file2 , folderLocationOfIndex)
        listOfIndexFiles = sorted(os.listdir(folderLocationOfIndex))
          
        



startTime = time.time()

parser = xml.sax.make_parser()
parser.setFeature(xml.sax.handler.feature_namespaces, 0)
Handler = Documenthandler()
parser.setContentHandler( Handler )
indexName = "index/"
os.mkdir("index")

directory = sys.argv[1] 
for filename in os.listdir(directory):
  dumpfile = directory + filename
  beforeParse = time.time() 
  parser.parse(dumpfile)
  print(dumpfile)
  print("Parse time : " + str(time.time() - beforeParse))


print("MergingStarted")
beforeMerge = time.time() 
mergeFiles(indexName)
print("Merge time : " + str(time.time() - beforeMerge))


print("splitting Started")
splitTime = time.time()
listOfIndexFiles = sorted(os.listdir(indexName))
if len(listOfIndexFiles) == 1:
  split(indexName + listOfIndexFiles[0],indexName)
  print("split time : " + str(time.time() - splitTime))
else:
  print("check split logic")

beforePickle = time.time()

statPath = 'titles/' +  'stat.txt'
st = open(statPath,'w')
st.write(str(docNo))
st.close()

dbfile = open(titlePickle, 'ab') 
pickle.dump(titleMap, dbfile)                      
dbfile.close() 


titleptr.close()

print("picle Save time : " + str(time.time() - beforePickle))

print("Total time : " + str(time.time() - startTime))