geonmin-kim's picture
Upload folder using huggingface_hub
d6585f5
#
# Pyserini: Reproducible IR research with sparse and dense representations
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
from pyserini.pyclass import autoclass
import json
import numpy as np
import pandas as pd
import spacy
import re
class Feature:
def name(self):
return self.extractor.getName()
class NormalizedTfIdf(Feature):
def __init__(self, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.NormalizedTfIdf')
self.extractor = Jclass(field, qfield)
class ProbalitySum(Feature):
def __init__(self, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.ProbalitySum')
self.extractor = Jclass(field, qfield)
class IbmModel1(Feature):
def __init__(self, path, field, tag, qfield):
Jclass = autoclass('io.anserini.ltr.feature.IbmModel1')
self.extractor = Jclass(path, field, tag, qfield)
class Proximity(Feature):
def __init__(self, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.Proximity')
self.extractor = Jclass(field, qfield)
class TpScore(Feature):
def __init__(self, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.TpScore')
self.extractor = Jclass(field, qfield)
class TpDist(Feature):
def __init__(self, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.TpDist')
self.extractor = Jclass(field, qfield)
class DocSize(Feature):
def __init__(self, field='contents'):
Jclass = autoclass('io.anserini.ltr.feature.DocSize')
self.extractor = Jclass(field)
class MatchingTermCount(Feature):
def __init__(self, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.MatchingTermCount')
self.extractor = Jclass(field, qfield)
class QueryLength(Feature):
def __init__(self, qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.QueryLength')
self.extractor = Jclass(qfield)
class SCS(Feature):
def __init__(self, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.SCS')
self.extractor = Jclass(field, qfield)
class SumMatchingTF(Feature):
def __init__(self, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.SumMatchingTF')
self.extractor = Jclass(field, qfield)
class QueryCoverageRatio(Feature):
def __init__(self, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.QueryCoverageRatio')
self.extractor = Jclass(field, qfield)
class RunList(Feature):
def __init__(self,filename,tag):
Jclass = autoclass('io.anserini.ltr.feature.RunList')
self.extractor = Jclass(filename,tag)
class UniqueTermCount(Feature):
def __init__(self, qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.UniqueTermCount')
self.extractor = Jclass(qfield)
class UnorderedSequentialPairs(Feature):
def __init__(self, gap=8, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.UnorderedSequentialPairs')
self.extractor = Jclass(gap, field, qfield)
class OrderedSequentialPairs(Feature):
def __init__(self, gap=8, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.OrderedSequentialPairs')
self.extractor = Jclass(gap, field, qfield)
class UnorderedQueryPairs(Feature):
def __init__(self, gap=8, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.UnorderedQueryPairs')
self.extractor = Jclass(gap, field, qfield)
class OrderedQueryPairs(Feature):
def __init__(self, gap=8, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.OrderedQueryPairs')
self.extractor = Jclass(gap, field, qfield)
class AvgPooler(Feature):
def __init__(self):
Jclass = autoclass('io.anserini.ltr.AvgPooler')
self.extractor = Jclass()
class SumPooler(Feature):
def __init__(self):
Jclass = autoclass('io.anserini.ltr.SumPooler')
self.extractor = Jclass()
class MedianPooler(Feature):
def __init__(self):
Jclass = autoclass('io.anserini.ltr.MedianPooler')
self.extractor = Jclass()
class MinPooler(Feature):
def __init__(self):
Jclass = autoclass('io.anserini.ltr.MinPooler')
self.extractor = Jclass()
class MaxPooler(Feature):
def __init__(self):
Jclass = autoclass('io.anserini.ltr.MaxPooler')
self.extractor = Jclass()
class VarPooler(Feature):
def __init__(self):
Jclass = autoclass('io.anserini.ltr.VarPooler')
self.extractor = Jclass()
class ConfidencePooler(Feature):
def __init__(self):
Jclass = autoclass('io.anserini.ltr.ConfidencePooler')
self.extractor = Jclass()
class MaxMinRatioPooler(Feature):
def __init__(self):
Jclass = autoclass('io.anserini.ltr.MaxMinRatioPooler')
self.extractor = Jclass()
class TfStat(Feature):
def __init__(self, pooler, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.TfStat')
self.extractor = Jclass(pooler.extractor, field, qfield)
class TfIdfStat(Feature):
def __init__(self, sublinear, pooler, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.TfIdfStat')
JBoolean = autoclass('java.lang.Boolean')
self.extractor = Jclass(JBoolean(sublinear), pooler.extractor, field, qfield)
class NormalizedTfStat(Feature):
def __init__(self, pooler, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.NormalizedTfStat')
self.extractor = Jclass(pooler.extractor, field, qfield)
class IdfStat(Feature):
def __init__(self, pooler, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.IdfStat')
self.extractor = Jclass(pooler.extractor, field, qfield)
class IcTfStat(Feature):
def __init__(self, pooler, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.IcTfStat')
self.extractor = Jclass(pooler.extractor, field, qfield)
class BM25Stat(Feature):
def __init__(self, pooler, k1=0.9, b=0.4, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.BM25Stat')
self.extractor = Jclass(pooler.extractor, k1, b, field, qfield)
class DfrInExpB2Stat(Feature):
def __init__(self, pooler, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.DfrInExpB2Stat')
self.extractor = Jclass(pooler.extractor, field, qfield)
class DphStat(Feature):
def __init__(self, pooler, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.DphStat')
self.extractor = Jclass(pooler.extractor, field, qfield)
class LmDirStat(Feature):
def __init__(self, pooler, mu=1000, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.LmDirStat')
self.extractor = Jclass(pooler.extractor, mu, field, qfield)
class DfrGl2Stat(Feature):
def __init__(self, pooler, field='contents', qfield='analyzed'):
Jclass = autoclass('io.anserini.ltr.feature.DfrGl2Stat')
self.extractor = Jclass(pooler.extractor, field, qfield)
class FeatureExtractor:
def __init__(self, index_dir, worker_num=1):
JFeatureExtractorUtils = autoclass('io.anserini.ltr.FeatureExtractorUtils')
self.utils = JFeatureExtractorUtils(index_dir, worker_num)
self.feature_name = []
def add(self, pyclass):
"""
add feature extractor; cannot add feature extractors in the middle of extraction
Parameters
----------
pyclass: Feature
an initialized feature extractor
"""
self.utils.add(pyclass.extractor)
self.feature_name.append(pyclass.name())
def feature_names(self):
"""
get all feature names
Returns
-------
List[str] all the feature names in order
"""
return self.feature_name
def lazy_extract(self, qid, doc_ids, query_dict):
input = {'qid': qid, 'docIds': doc_ids}
input.update(query_dict)
self.utils.lazyExtract(json.dumps(input))
def batch_extract(self, tasks):
need_rows = 0
for task in tasks:
self.lazy_extract(task['qid'], task['docIds'], task['query_dict'])
need_rows += len(task['docIds'])
feature_name = self.feature_names()
feature = np.zeros(shape=(need_rows, len(feature_name)), dtype=np.float32)
idx = 0
for task in tasks:
flattened = self.get_result(task['qid'])
feature[idx:idx+len(task['docIds']),:] = flattened.reshape(len(task['docIds']), len(feature_name))
idx += len(task['docIds'])
return pd.DataFrame(feature, columns=feature_name)
def get_result(self, qid):
res = self.utils.getResult(qid).tostring()
dt = np.dtype(np.float32)
dt = dt.newbyteorder('>')
return np.frombuffer(res, dt)
class SpacyTextParser:
def __init__(self, model_name,
remove_punct=True,
sent_split=False,
keep_only_alpha_num=False,
lower_case=True,
enable_POS=True):
disable_list = ['ner', 'parser']
if not enable_POS:
disable_list.append('tagger')
print('Disabled Spacy components: ', disable_list)
self._nlp = spacy.load(model_name, disable=disable_list)
if sent_split:
sentencizer = self._nlp.create_pipe("sentencizer")
self._nlp.add_pipe(sentencizer)
self._remove_punct = remove_punct
sw = ['a', 'about', 'above', 'according', 'across', 'after',
'afterwards', 'again', 'against', 'albeit', 'all', 'almost',
'alone', 'along', 'already', 'also', 'although', 'always', 'am',
'among', 'amongst', 'an', 'and', 'another', 'any', 'anybody', 'anyhow',
'anyone', 'anything', 'anyway', 'anywhere', 'apart', 'are', 'around',
'as', 'at', 'av', 'be', 'became', 'because', 'become', 'becomes',
'becoming', 'been', 'before', 'beforehand', 'behind', 'being', 'below',
'beside', 'besides', 'between', 'beyond', 'both', 'but', 'by', 'can',
'cannot', 'canst', 'certain', 'cf', 'choose', 'contrariwise', 'cos',
'could', 'cu', 'day', 'do', 'does', "doesn't", 'doing', 'dost', 'doth',
'double', 'down', 'dual', 'during', 'each', 'either', 'else', 'elsewhere',
'enough', 'et', 'etc', 'even', 'ever', 'every', 'everybody', 'everyone',
'everything', 'everywhere', 'except', 'excepted', 'excepting', 'exception',
'exclude', 'excluding', 'exclusive', 'far', 'farther', 'farthest', 'few',
'ff', 'first', 'for', 'formerly', 'forth', 'forward', 'from', 'front',
'further', 'furthermore', 'furthest', 'get', 'go', 'had', 'halves', 'hardly',
'has', 'hast', 'hath', 'have', 'he', 'hence', 'henceforth', 'her', 'here',
'hereabouts', 'hereafter', 'hereby', 'herein', 'hereto', 'hereupon', 'hers',
'herself', 'him', 'himself', 'hindmost', 'his', 'hither', 'hitherto', 'how',
'however', 'howsoever', 'i', 'ie', 'if', 'in', 'inasmuch', 'inc', 'include',
'included', 'including', 'indeed', 'indoors', 'inside', 'insomuch', 'instead',
'into', 'inward', 'inwards', 'is', 'it', 'its', 'itself', 'just', 'kind', 'kg',
'km', 'last', 'latter', 'latterly', 'less', 'lest', 'let', 'like', 'little', 'ltd',
'many', 'may', 'maybe', 'me', 'meantime', 'meanwhile', 'might', 'moreover', 'most',
'mostly', 'more', 'mr', 'mrs', 'ms', 'much', 'must', 'my', 'myself', 'namely', 'need',
'neither', 'never', 'nevertheless', 'next', 'no', 'nobody', 'none', 'nonetheless',
'noone', 'nope', 'nor', 'not', 'nothing', 'notwithstanding', 'now', 'nowadays',
'nowhere', 'of', 'off', 'often', 'ok', 'on', 'once', 'one', 'only', 'onto', 'or',
'other', 'others', 'otherwise', 'ought', 'our', 'ours', 'ourselves', 'out', 'outside',
'over', 'own', 'per', 'perhaps', 'plenty', 'provide', 'quite', 'rather', 'really',
'round', 'said', 'sake', 'same', 'sang', 'save', 'saw', 'see', 'seeing', 'seem', 'seemed',
'seeming', 'seems', 'seen', 'seldom', 'selves', 'sent', 'several', 'shalt', 'she', 'should',
'shown', 'sideways', 'since', 'slept', 'slew', 'slung', 'slunk', 'smote', 'so', 'some',
'somebody', 'somehow', 'someone', 'something', 'sometime', 'sometimes', 'somewhat', 'somewhere',
'spake', 'spat', 'spoke', 'spoken', 'sprang', 'sprung', 'stave', 'staves', 'still', 'such',
'supposing', 'than', 'that', 'the', 'thee', 'their', 'them', 'themselves', 'then', 'thence',
'thenceforth', 'there', 'thereabout', 'thereabouts', 'thereafter', 'thereby', 'therefore',
'therein', 'thereof', 'thereon', 'thereto', 'thereupon', 'these', 'they', 'this', 'those',
'thou', 'though', 'thrice', 'through', 'throughout', 'thru', 'thus', 'thy', 'thyself', 'till',
'to', 'together', 'too', 'toward', 'towards', 'ugh', 'unable', 'under', 'underneath', 'unless',
'unlike', 'until', 'up', 'upon', 'upward', 'upwards', 'us', 'use', 'used', 'using', 'very', 'via',
'vs', 'want', 'was', 'we', 'week', 'well', 'were', 'what', 'whatever', 'whatsoever', 'when',
'whence', 'whenever', 'whensoever', 'where', 'whereabouts', 'whereafter', 'whereas', 'whereat',
'whereby', 'wherefore', 'wherefrom', 'wherein', 'whereinto', 'whereof', 'whereon', 'wheresoever',
'whereto', 'whereunto', 'whereupon', 'wherever', 'wherewith', 'whether', 'whew', 'which',
'whichever', 'whichsoever', 'while', 'whilst', 'whither', 'who', 'whoa', 'whoever', 'whole',
'whom', 'whomever', 'whomsoever', 'whose', 'whosoever', 'why', 'will', 'wilt', 'with', 'within',
'without', 'worse', 'worst', 'would', 'wow', 'ye', 'yet', 'year', 'yippee', 'you', 'your', 'yours',
'yourself', 'yourselves', "n't", "'d", "'ll", "'m", "'re", "'s", "'ves"]
stopwords = set(sw)
self._stopwords = frozenset([w.lower() for w in stopwords])
self._keep_only_alpha_num = keep_only_alpha_num
self._lower_case = lower_case
@staticmethod
def _basic_clean(text):
return text.replace("’", "'")
def __call__(self, text):
return self._nlp(SpacyTextParser._basic_clean(text))
def is_alpha_num(self, s):
return s and (re.match("^[a-zA-Z-_.0-9]+$", s) is not None)
def proc_text(self, text):
lemmas = []
tokens = []
doc = self(text)
for tokObj in doc:
if self._remove_punct and tokObj.is_punct:
continue
lemma = tokObj.lemma_
text = tokObj.text
if self._keep_only_alpha_num and not self.is_alpha_num(text):
continue
tok1 = text.lower()
tok2 = lemma.lower()
if tok1 in self._stopwords or tok2 in self._stopwords:
continue
if self._lower_case:
text = text.lower()
lemma = lemma.lower()
lemmas.append(lemma)
tokens.append(text)
return ' '.join(lemmas), ' '.join(tokens)