mirror of
https://github.com/wassname/Castor.git
synced 2026-08-20 12:00:37 +08:00
Michael Tu
f7a0167b81
* Refactor main README * Update Anserini Dependency docs * Update idf baseline and Kim CNN docs to use Castor-data * Update remaining READMEs to reference Castor-data * Change default path from data to Castor-data * Fix wrong order of embeddings path
188 lines
7.6 KiB
Python
188 lines
7.6 KiB
Python
import argparse
|
|
import os
|
|
import sys
|
|
import re
|
|
import numpy as np
|
|
from collections import defaultdict
|
|
import string
|
|
import subprocess
|
|
import shlex
|
|
|
|
import nltk
|
|
nltk.download('stopwords', quiet=True)
|
|
|
|
from nltk.stem.porter import PorterStemmer
|
|
from nltk.corpus import stopwords
|
|
|
|
def read_in_data(datapath, set_name, file, stop_and_stem=False, stop_punct=False, dash_split=False):
|
|
data = []
|
|
with open(os.path.join(datapath, set_name, file)) as inf:
|
|
data = [line.strip() for line in inf.readlines()]
|
|
|
|
if dash_split:
|
|
def split_hyphenated_words(sentence):
|
|
rtokens = []
|
|
for term in sentence.split():
|
|
for t in term.split('-'):
|
|
if t:
|
|
rtokens.append(t)
|
|
return ' '.join(rtokens)
|
|
data = [split_hyphenated_words(sentence) for sentence in data]
|
|
|
|
if stop_punct:
|
|
regex = re.compile('[{}]'.format(re.escape(string.punctuation)))
|
|
def remove_punctuation(sentence):
|
|
rtokens = []
|
|
for term in sentence.split():
|
|
for t in regex.sub(' ', term).strip().split():
|
|
if t:
|
|
rtokens.append(t)
|
|
return ' '.join(rtokens)
|
|
data = [remove_punctuation(sentence) for sentence in data]
|
|
|
|
if stop_and_stem:
|
|
stemmer = PorterStemmer()
|
|
stoplist = set(stopwords.words('english'))
|
|
def stop_stem(sentence):
|
|
return ' '.join([stemmer.stem(word) for word in sentence.split() \
|
|
if word not in stoplist])
|
|
data = [stop_stem(sentence) for sentence in data]
|
|
return data
|
|
|
|
|
|
def compute_idfs(data, dash_split=False):
|
|
term_idfs = defaultdict(float)
|
|
for doc in list(data):
|
|
for term in list(set(doc.split())):
|
|
if dash_split:
|
|
assert('-' not in term)
|
|
term_idfs[term] += 1.0
|
|
N = len(data)
|
|
for term, n_t in term_idfs.items():
|
|
term_idfs[term] = np.log(N/(1+n_t))
|
|
return term_idfs
|
|
|
|
def fetch_idfs_from_index(data, dash_split, indexPath):
|
|
regex = re.compile('[{}]'.format(re.escape(string.punctuation)))
|
|
term_idfs = defaultdict(float)
|
|
all_terms = set([term for doc in list(data) for term in doc.split()])
|
|
with open('dataset.vocab', 'w') as vf:
|
|
for term in list(all_terms):
|
|
if dash_split:
|
|
assert('-' not in term)
|
|
print(term, file=vf)
|
|
|
|
fetchIDF_cmd = \
|
|
"sh ../idf_baseline/target/appassembler/bin/FetchTermIDF -index {} -vocabFile {}".\
|
|
format(indexPath, 'dataset.vocab')
|
|
pargs = shlex.split(fetchIDF_cmd)
|
|
p = subprocess.Popen(pargs, stdout=subprocess.PIPE, stderr=subprocess.PIPE, \
|
|
bufsize=1, universal_newlines=True)
|
|
pout, perr = p.communicate()
|
|
|
|
lines = str(pout).split('\n')
|
|
for line in lines:
|
|
if not line:
|
|
continue
|
|
fields = line.strip().split("\t")
|
|
term, weight = fields[0], fields[-1]
|
|
term_idfs[term] = float(weight)
|
|
|
|
for line in str(perr).split('\n'):
|
|
print('Warning: '+line)
|
|
return term_idfs
|
|
|
|
def compute_idf_sum_similarity(questions, answers, term_idfs):
|
|
# compute IDF sums for common_terms
|
|
idf_sum_similarity = np.zeros(len(questions))
|
|
for i in range(len(questions)):
|
|
q = questions[i]
|
|
a = answers[i]
|
|
q_terms = set(q.split())
|
|
a_terms = set(a.split())
|
|
common_terms = q_terms.intersection(a_terms)
|
|
idf_sum_similarity[i] = np.sum([term_idfs[term] for term in list(common_terms)])
|
|
|
|
return idf_sum_similarity
|
|
|
|
|
|
def write_out_idf_sum_similarities(qids, questions, answers, term_idfs, outfile, dataset):
|
|
with open(outfile, 'w') as outf:
|
|
idf_sum_similarity = compute_idf_sum_similarity(questions, answers, term_idfs)
|
|
old_qid = 0
|
|
docid_c = 0
|
|
for i in range(len(questions)):
|
|
if qids[i] != old_qid and dataset.endswith('WikiQA'):
|
|
docid_c = 0
|
|
old_qid = qids[i]
|
|
print('{} 0 {} 0 {} idfbaseline'.format(qids[i], docid_c,
|
|
idf_sum_similarity[i]),
|
|
file=outf)
|
|
docid_c += 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
ap = argparse.ArgumentParser(description="uses idf weights from the question-answer pairs only,\
|
|
and not from the whole corpus")
|
|
ap.add_argument('qa_data', help="path to the QA dataset",
|
|
choices=['../../Castor-data/TrecQA', '../../Castor-data/WikiQA'])
|
|
ap.add_argument('outfile_prefix', help="output file prefix")
|
|
ap.add_argument('--ignore-test', help="does not consider test data when computing IDF of terms",
|
|
action="store_true")
|
|
ap.add_argument("--stop-and-stem", help='performs stopping and stemming', action="store_true")
|
|
ap.add_argument("--stop-punct", help='removes punctuation', action="store_true")
|
|
ap.add_argument("--dash-split", help="split words containing hyphens", action="store_true")
|
|
ap.add_argument("--index-for-corpusIDF", help="fetches idf from Index. provide index path. will\
|
|
generate a vocabFile")
|
|
|
|
args = ap.parse_args()
|
|
|
|
# read in the data
|
|
train_data, dev_data, test_data = 'train', 'dev', 'test'
|
|
if args.qa_data.endswith('TrecQA'):
|
|
train_data, dev_data, test_data = 'train-all', 'raw-dev', 'raw-test'
|
|
|
|
train_que = read_in_data(args.qa_data, train_data, 'a.toks',
|
|
args.stop_and_stem, args.stop_punct, args.dash_split)
|
|
train_ans = read_in_data(args.qa_data, train_data, 'b.toks',
|
|
args.stop_and_stem, args.stop_punct, args.dash_split)
|
|
|
|
dev_que = read_in_data(args.qa_data, dev_data, 'a.toks',
|
|
args.stop_and_stem, args.stop_punct, args.dash_split)
|
|
dev_ans = read_in_data(args.qa_data, dev_data, 'b.toks',
|
|
args.stop_and_stem, args.stop_punct, args.dash_split)
|
|
|
|
test_que = read_in_data(args.qa_data, test_data, 'a.toks',
|
|
args.stop_and_stem, args.stop_punct, args.dash_split)
|
|
test_ans = read_in_data(args.qa_data, test_data, 'b.toks',
|
|
args.stop_and_stem, args.stop_punct, args.dash_split)
|
|
|
|
all_data = train_que + dev_que + train_ans + dev_ans
|
|
|
|
if not args.ignore_test:
|
|
all_data += test_ans
|
|
all_data += test_que
|
|
|
|
# compute inverse document frequencies for terms
|
|
if not args.index_for_corpusIDF:
|
|
term_idfs = compute_idfs(set(all_data), args.dash_split)
|
|
else:
|
|
term_idfs = fetch_idfs_from_index(set(all_data), args.dash_split, args.index_for_corpusIDF)
|
|
|
|
# write out in trec_eval format
|
|
write_out_idf_sum_similarities(read_in_data(args.qa_data, train_data, 'id.txt'),
|
|
train_que, train_ans, term_idfs,
|
|
'{}.{}.idfsim'.format(args.outfile_prefix, train_data),
|
|
args.qa_data)
|
|
|
|
write_out_idf_sum_similarities(read_in_data(args.qa_data, dev_data, 'id.txt'),
|
|
dev_que, dev_ans, term_idfs,
|
|
'{}.{}.idfsim'.format(args.outfile_prefix, dev_data),
|
|
args.qa_data)
|
|
|
|
write_out_idf_sum_similarities(read_in_data(args.qa_data, test_data, 'id.txt'),
|
|
test_que, test_ans, term_idfs,
|
|
'{}.{}.idfsim'.format(args.outfile_prefix, test_data),
|
|
args.qa_data)
|
|
|