retriever.py

"""Author: Zhentao Huang (Github link for this assignment: https://github.com/ZhentaoHuang/CIS-6190-Assignment-2)
This file is used to Perform the Online Processing task. Type -h for help.
"""

from nltk.corpus import stopwords
from nltk.stem import PorterStemmer
from bisect import bisect_left

import argparse
import re
import string
import math





def read_files(dictionary_file, postings_file, docids_file):
    """This function is used to read the three files generated by indexer.py

    Args:
        dictionary_file (file): input file dictionary.txt
        postings_file (file): input file postings.txt
        docids_file (file): input file docids.txt

    Returns:
        list: three list contains the contents from the files
    """
    # Construct keyword list
    dictionary = []
    lines = dictionary_file.readlines()  # Read the dictionary file
    count = 0
    for line in lines:
        stem, df = line[:-1].split(" ")
        dictionary.append((stem, count))  # Convert the document frequency to offset
        count = count + int(df)  # Adjust the offset for next stem
    
    # Construct posting list
    postings = []
    lines = postings_file.readlines()
    for line in lines:
        did, tf = line[:-1].split(" ")
        postings.append((int(did), int(tf)))
    
    # Construct docid lsit
    docids = []
    lines = docids_file.readlines()
    for line in lines:
        words = line[:-1].split(" ")
        docid = words[0]
        start_line_number = words[-1]
        title = " ".join(words[1:-1])
        docids.append((docid, title, start_line_number))
    
    
    return dictionary, postings, docids

def preproc(query):
    """This function is used to perform the preprocessing task. The code is taken from preprocessor.py

    Args:
        query (string): command line query

    Returns:
        List of preprocessed words
    """
    words = query.split(" ")
    ps = PorterStemmer()
    numbers = re.compile(r"[+|-]?\d+(\.\d+)?") # Integers and real numbers, with possible positive and negative signs
    word_list= []
    for word in words:
                match = numbers.match(word)  # Remove all the numbers
                if match is None and word not in string.punctuation:    # Remove the punctuation
                    if word not in stopwords.words("english"):  # Remove the stopwords

                        word_list.append(ps.stem(word.lower()))    # Lowercase and Stemming

    
    return word_list

def BinSearch(nums, x):
    """Binary Search taken from the sample program

    Args:
        nums (list): list for binary search
        x (string): item to search

    Returns:
        int: index
    """
    i = bisect_left(nums, x)
    if i != len(nums) and nums[i] == x:
        return i
    else:
        return -1

def get_df(index, dictionary, postings):
    """Calculate the document frequency given the lists

    Args:
        index (int): the index of stem in the dictionary
        dictionary (list): the dictionary list
        postings (list): the postings list

    Returns:
        int: document frequency
    """
    
    # index = -1 if not found
    if (index >= 0):
        # if the index pointing to the last doc
        if(index == len(dictionary) - 1):
            df = len(postings) - dictionary[index]
        else:
            df = dictionary[index + 1][1] - dictionary[index][1]
    else:
        df = 0

    return df

def online_process(dictionary, postings, docids):
    """This function is the main function to perform the online_process task. It allows user to iteratively search for related documents.

    Args:
        dictionary (file): dictionary.txt
        postings (file): postings.txt
        docids (file): docids.txt
    """

    n = len(docids)    # Total number of documents
    print(str(n) + " documents were loded.")
    query = " "
    keys = [r[0] for r in dictionary]   # Construct a list for binary search
    while(True):
        query = input("Please enter for quering (\"quit\" or \"q\" to quit): ")
        if(query == "quit" or query == "q"):    # If the user input quit or q then end it
            break
        print("You entered: " + query)
        words = preproc(query)     # Preprocess the query
        weights = []    # The list used to store the weights/similarities and did: [(did, weight)]
        for i in range(len(docids)):
            weights.append((i, 0))  # Initialize with zeros


        for word in words:
            index = BinSearch(keys, word)   # Search for the stem
            if (index < 0):
                continue    # Not found
            df = get_df(index, dictionary, postings)
 
            # Get the offset and convert to start and end
            if(index == len(keys) - 1):     # Situation for last one offset
                start_offset = dictionary[index][1]
                end_offset = len(postings)
            else:
                start_offset = dictionary[index][1]
                end_offset = dictionary[index+1][1]
            
            # Calculate the weights based on occured document
            for j in range(start_offset, end_offset):
               
                did = postings[j][0]
                tf = postings[j][1]
                doc_w = tf * math.log2(n/df)
                q_w = 1 * math.log2(n/df)
               
                weights[did] = (did, weights[did][1] + doc_w * q_w)
     


        # Sort the weights in a descending order to find the top-10 match
        weights.sort(key=lambda tup: tup[1], reverse=True)
        count = 0
        for i in range(10):
            
            
            inner_product = weights[i][1]

            if(inner_product > 0):
                id = docids[weights[i][0]][0]
                title = docids[weights[i][0]][1]
                print(str(i + 1) + ". Title: " + title + "\tDocid: " +id + "\tSimilarity:" + str(inner_product))
                count = count + 1
   
        # Print the final output
        if(count == 0):
            print("No document found.")
        elif(count < 10):
            print("Only " + str(count) + " document found.")
        else:
            print("Top-10 result displayed.")
    
def arg_parse():
    """This function is used for command line argument parsing. It utilizes the argparse library. It provides the user with file choice at runtime.

    Returns:
        args: command line arguments
    """

    # Initialize parser
    parser = argparse.ArgumentParser()
    
    # Adding optional argument
    parser.add_argument("-d", "--Dictionary", default="dictionary.txt", help = "The input dictionary file (default: dictionary.txt)")
    parser.add_argument("-p", "--Postings", default="postings.txt", help = "The input postings file (default: postings.txt)")
    parser.add_argument("-s", "--Docids", default="docids.txt", help = "The input docids file (default: docids.txt)")

    
    # Read arguments from command line
    args = parser.parse_args()
    

    print("Output as: % s" % args.Dictionary)
    print("Output as: % s" % args.Postings)
    print("Output as: % s" % args.Docids)

    return args

def main(args):

    #input = open(args.Input, 'r')
    dictionary_file = open(args.Dictionary, 'r')
    postings_file = open(args.Postings, 'r')
    docids_file = open(args.Docids, 'r')

    
    dictionary, postings, docids = read_files(dictionary_file, postings_file, docids_file)
    online_process(dictionary, postings, docids)
    
    
    dictionary_file.close()
    postings_file.close()
    docids_file.close()

if __name__ == "__main__":

    args = arg_parse()
    main(args)