54 lines
1.5 KiB
Python
54 lines
1.5 KiB
Python
import re
|
|
import pandas as pd
|
|
import nltk
|
|
import os
|
|
import argparse
|
|
|
|
def read_document(path):
|
|
'''Reads a file when provided with its path, and returns a string
|
|
containing the lines of the file.'''
|
|
file_given = open(path)
|
|
f = ""
|
|
for line in file_given:
|
|
f += line + " "
|
|
file_given.close()
|
|
return f
|
|
|
|
def apply_preprocessing(f):
|
|
'''Removes non-alphabetic characters, replaces all whitespace characters
|
|
with a single whitespace, and changes all uppercase characters to
|
|
lowercase'''
|
|
f = re.sub(r'[^a-zA-Z\s]', r'', f)
|
|
f = re.sub(r'\s+', r' ', f)
|
|
f = f.lower()
|
|
return f
|
|
|
|
# parse input arguments
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('keywords', nargs = '+', help = 'keywords to search for \
|
|
(1-5 keywords accepted)')
|
|
args = parser.parse_args()
|
|
if len(args.keywords) > 5:
|
|
print("Too many keywords.")
|
|
quit()
|
|
|
|
# load document IDs from csv
|
|
df = pd.read_csv('partb1.csv', encoding = 'ISO-8859-1')
|
|
doc_ids = pd.Series(data = df.documentID.tolist(), \
|
|
index = df.filename.tolist())
|
|
documents = doc_ids.index.tolist()
|
|
matched_doc_ids = []
|
|
|
|
os.chdir(os.getcwd() + '/cricket')
|
|
|
|
# search through each document for the keywords
|
|
for doc in documents:
|
|
f = read_document(doc)
|
|
f = apply_preprocessing(f)
|
|
|
|
tokens = nltk.word_tokenize(f)
|
|
# only add the document ID if all the keywords are in the token list
|
|
if all(keyword in tokens for keyword in args.keywords):
|
|
matched_doc_ids.append(doc_ids.get(doc))
|
|
|
|
print(matched_doc_ids)
|