def test_search_document_id():
md = [{
'file_path': '/test',
'document_id': 2
}, {
'file_path': '/test2',
'document_id': 1
}, {
'file_path': '/test3',
'document_id': 7
}, {
'file_path': '/test8',
'document_id': 9
}, {
'file_path': '/test9',
'document_id': 4
}]
for idx, el in enumerate(md):
el['internal_id'] = idx
dbi = DocumentIndex.from_list(md)
dbi._make_relative_paths()
query = pd.DataFrame([{
'internal_id': 1
}, {
'internal_id': 2
}, {
'internal_id': 1
}])
sres = dbi.search(query)
assert_equal(sres.internal_id.values, [1, 2, 1])
assert_array_equal(sorted(sres.columns),
sorted(['internal_id', 'file_path', 'document_id']))
# make sure we use internal id first
query = pd.DataFrame([{
'internal_id': 1,
'document_id': 2
}, {
'internal_id': 2,
'document_id': 2
}, {
'internal_id': 1,
'document_id': 2
}])
sres = dbi.search(query)
assert_equal(sres.internal_id.values, [1, 2, 1])
query = pd.DataFrame([{
'document_id': 4
}, {
'document_id': 9
}, {
'document_id': 2
}])
sres = dbi.search(query)
assert_equal(sres.internal_id.values, [4, 3, 0])
def test_ingestion_metadata(n_fields):
metadata = []
for idx, fname in enumerate(fnames_in_abs):
el = {'file_path': fname}
if n_fields >= 2:
el['document_id'] = 'a' + str(idx + 100)
if n_fields >= 3:
el['rendition_id'] = 1
metadata.append(el)
dbi = DocumentIndex.from_list(metadata)
dbi._make_relative_paths()
data_dir_res, filenames, db = dbi.data_dir, dbi.filenames_, dbi.data
if n_fields in [1, 2]:
columns_ref = sorted(['file_path', 'document_id', 'internal_id'])
elif n_fields == 3:
columns_ref = sorted(
['file_path', 'document_id', 'rendition_id', 'internal_id'])
assert_array_equal(sorted(db.columns.values), columns_ref)
assert_array_equal(
[os.path.normpath(os.path.join(data_dir_res, el)) for el in filenames],
[os.path.join(data_dir_res, el) for el in db.file_path.values])
def test_bad_search_document_rendition_id():
md = [{
'file_path': '/test',
'document_id': 0,
'rendition_id': 0
}, {
'file_path': '/test2',
'document_id': 0,
'rendition_id': 1
}, {
'file_path': '/test3',
'document_id': 1,
'rendition_id': 0
}, {
'file_path': '/test8',
'document_id': 2,
'rendition_id': 0
}, {
'file_path': '/test9',
'document_id': 3,
'rendition_id': 0
}]
for idx, el in enumerate(md):
el['internal_id'] = idx
# can always index with internal_id
dbi = DocumentIndex.from_list(md)
query = pd.DataFrame([{
'internal_id': 1
}, {
'internal_id': 2
}, {
'document_id': 1
}])
with pytest.raises(NotFound):
sres = dbi.search(query)
def append(self, dataset_definition, data_dir=None):
""" Add some documents to the dataset
This is by no mean an efficient operation, processing all the files
at once might be more suitable in most occastions.
"""
from freediscovery.engine.lsi import _LSIWrapper
dsid_dir = self.dsid_dir
db_old = self.db_.data
internal_id_offset = db_old.internal_id.max()
db_extra = DocumentIndex.from_list(dataset_definition, data_dir,
internal_id_offset + 1, dsid_dir)
db_new = db_extra.data
vect = self.vect_
tfidf = self.tfidf_
filenames_new = list(db_new.file_path.values)
# write down the new features file
X_new_raw = vect.transform(filenames_new)
X_new = tfidf.transform(X_new_raw)
X_old = self._load_features()
X = scipy.sparse.vstack((X_new, X_old))
joblib.dump(X, str(dsid_dir / 'features'))
# write down the new filenames file
filenames_old = list(self.filenames_)
filenames = filenames_old + filenames_new
data_dir = DocumentIndex._detect_data_dir(filenames)
self._pars['data_dir'] = data_dir
self._filenames = [os.path.relpath(el, data_dir)
for el in filenames]
with (dsid_dir / 'filenames').open('wb') as fh:
pickle.dump(self._filenames, fh)
del db_new['file_path']
# write down the new pars file
self._pars = self.pars_
self._pars['n_samples'] = len(filenames)
with (dsid_dir / 'pars').open('wb') as fh:
pickle.dump(self._pars, fh)
# write down the new database file
db = pd.concat((db_old, db_new))
if 'file_path' in db.columns:
del db['file_path']
db.to_pickle(str(dsid_dir / 'db'))
self._db = DocumentIndex(self.pars_['data_dir'], db)
# find all exisisting LSI models and update them as well
if (dsid_dir / 'lsi').exists():
for lsi_id in os.listdir(str(dsid_dir / 'lsi')):
lsi_obj = _LSIWrapper(cache_dir=self.cache_dir,
mid=lsi_id)
lsi_obj.append(X_new)
# remove all trained models for this dataset
for model_type in ['categorizer', 'dupdet', 'cluster', 'threading']:
if (dsid_dir / model_type).exists():
for mid in os.listdir(str(dsid_dir / model_type)):
shutil.rmtree(str(dsid_dir / model_type / mid))
def ingest(self, data_dir=None, file_pattern='.*', dir_pattern='.*',
dataset_definition=None, vectorize=True,
document_id_generator='indexed_file_path',
):
"""Perform data ingestion
Parameters
----------
data_dir : str
path to the data directory (used only if metadata not provided),
default: None
dataset_defintion : list of dicts
a list of dictionaries with keys
['file_path', 'document_id', 'rendition_id']
describing the data ingestion (this overwrites data_dir)
vectorize : bool (default: True)
"""
dsid_dir = self.cache_dir / self.dsid
if (dsid_dir / 'db').exists():
raise ValueError('Dataset {} already vectorized!'
.format(self.dsid))
db_list = list(sorted(dsid_dir.glob('db*')))
if len(db_list) == 0:
internal_id_offset = -1
elif len(db_list) >= 1:
internal_id_offset = int(db_list[-1].name[3:])
pars = self.pars_
if pars.get('column_ids', None) is not None:
if dataset_definition is None:
raise ValueError("CSV files can only be privided using "
"`dataset_definition` parameter")
else:
if len(dataset_definition) > 1:
raise ValueError(
"Only one CSV can be provided at a time"
)
file_path = dataset_definition[0]['file_path']
X = pd.read_csv(
file_path, sep=pars['column_separator'], header=None)
dataset_definition = [
{'file_path': f"{file_path}:{idx}", 'document_id': idx}
for idx in range(len(X))]
db = DocumentIndex.from_list(
dataset_definition, data_dir,
internal_id_offset + 1, dsid_dir,
document_id_generator=document_id_generator)
elif dataset_definition is not None:
db = DocumentIndex.from_list(
dataset_definition, data_dir,
internal_id_offset + 1, dsid_dir,
document_id_generator=document_id_generator)
elif data_dir is not None:
db = DocumentIndex.from_folder(
data_dir, file_pattern, dir_pattern,
internal_id_offset + 1,
document_id_generator=document_id_generator)
else:
db = None
if db is not None:
data_dir = db.data_dir
batch_suffix = '.{:09}'.format(db.data.internal_id.iloc[-1])
self._filenames = db.data.file_path.values.tolist()
del db.data['file_path']
if 'file_path' in db.data.columns:
del db.data['file_path']
db.data.to_pickle(str(dsid_dir / ('db' + batch_suffix)))
with (dsid_dir / ('filenames' + batch_suffix)).open('wb') as fh:
pickle.dump(self._filenames, fh)
self._db = db
if vectorize:
db_list = list(sorted(dsid_dir.glob('db*')))
filenames_list = list(sorted(dsid_dir.glob('filenames*')))
if len(db_list) == 0:
raise ValueError('No ingested files found!')
if len(db_list) == 1:
with filenames_list[0].open('rb') as fh:
filenames_concat = pickle.load(fh)
elif len(db_list) >= 2:
# accumulate different batches into a single file
# filename file
filenames_concat = []
for fname in filenames_list:
with fname.open('rb') as fh:
filenames_concat += pickle.load(fh)
if self.pars_['data_dir'] is None:
data_dir = DocumentIndex._detect_data_dir(filenames_concat)
self._pars['data_dir'] = data_dir
else:
data_dir = self._pars['data_dir']
self._filenames = [os.path.relpath(el, data_dir)
for el in filenames_concat]
with (dsid_dir / 'filenames').open('wb') as fh:
pickle.dump(self._filenames, fh)
for fname in filenames_list:
fname.unlink()
# save databases
if len(db_list) == 1:
db_list[0].rename(dsid_dir / 'db')
self.db_.filenames_ = self._filenames
self.db_.data['file_path'] = self._filenames
elif len(db_list) >= 2:
db_concat = []
for fname in db_list:
db_concat.append(pd.read_pickle(str(fname)))
db_new = pd.concat(db_concat, axis=0)
db_new.filenames_ = self._filenames
db_new.set_index('internal_id', drop=False, inplace=True)
self._db = DocumentIndex(data_dir, db_new)
if 'file_path' in db_new.columns:
del db_new['file_path']
db_new.to_pickle(str(dsid_dir / 'db'))
# save parameters
self._pars['n_samples'] = len(self._filenames)
self._pars['data_dir'] = data_dir
with (dsid_dir / 'pars').open('wb') as fh:
pickle.dump(self._pars, fh)
self.transform()
if (dsid_dir / 'raw').exists():
shutil.rmtree(str(dsid_dir / 'raw'))
if db is None and not vectorize:
raise ValueError('At least one of data_dir, dataset_definition, '
'vectorize parameters must be provided!')
return
def test_search_document_rendition_id():
md = [{
'file_path': '/test',
'document_id': 0,
'rendition_id': 0
}, {
'file_path': '/test2',
'document_id': 0,
'rendition_id': 1
}, {
'file_path': '/test3',
'document_id': 1,
'rendition_id': 0
}, {
'file_path': '/test8',
'document_id': 2,
'rendition_id': 0
}, {
'file_path': '/test9',
'document_id': 3,
'rendition_id': 0
}]
for idx, el in enumerate(md):
el['internal_id'] = idx
# can always index with internal_id
dbi = DocumentIndex.from_list(md)
dbi._make_relative_paths()
query = pd.DataFrame([{
'internal_id': 1
}, {
'internal_id': 2
}, {
'internal_id': 1
}])
sres = dbi.search(query)
assert_equal(sres.internal_id.values, [1, 2, 1])
assert_array_equal(
sorted(sres.columns),
sorted(['internal_id', 'file_path', 'document_id', 'rendition_id']))
# the internal id is not sufficient to fully index documents in this case
query = pd.DataFrame([{
'document_id': 0
}, {
'document_id': 1
}, {
'document_id': 2
}])
with pytest.raises(ValueError):
sres = dbi.search(query)
query = pd.DataFrame([{
'document_id': 0,
'rendition_id': 0
}, {
'document_id': 1,
'rendition_id': 0
}, {
'document_id': 2,
'rendition_id': 0
}])
sres = dbi.search(query)
assert_equal(sres.internal_id.values, [0, 2, 3])
def test_ingestion_render(return_file_path):
def _process_results(rd):
rd = pd.DataFrame(rd)
if return_file_path:
assert 'file_path' in rd.columns
del rd['file_path']
return rd
# make it a binary variable
return_file_path = (return_file_path == 'return_file_path')
md = [{
'file_path': '/test',
'document_id': 2
}, {
'file_path': '/test2',
'document_id': 1
}, {
'file_path': '/test3',
'document_id': 7
}, {
'file_path': '/test8',
'document_id': 9
}, {
'file_path': '/test9',
'document_id': 4
}]
for idx, el in enumerate(md):
el['internal_id'] = idx
dbi = DocumentIndex.from_list(md)
query = pd.DataFrame([{
'a': 2,
'internal_id': 3
}, {
'a': 4,
'internal_id': 1
}])
res = pd.DataFrame([{
'a': 2,
'internal_id': 3,
'document_id': 9
}, {
'a': 4,
'internal_id': 1,
'document_id': 1
}])
rd = dbi.render_dict(query, return_file_path=return_file_path)
rd = _process_results(rd)
assert_frame_equal(rd, res)
rd = dbi.render_dict(return_file_path=return_file_path)
rd = _process_results(rd)
assert_frame_equal(rd.loc[[0]],
pd.DataFrame([{
'internal_id': 0,
'document_id': 2
}]))
assert len(rd) == len(md)
rd = dbi.render_list(res, return_file_path=return_file_path)
rd = _process_results(rd)
assert sorted(rd.keys()) == sorted(['internal_id', 'document_id', 'a'])
assert_frame_equal(pd.DataFrame(rd),
pd.DataFrame([{
'a': 2,
'internal_id': 3,
'document_id': 9
}, {
'a': 4,
'internal_id': 1,
'document_id': 1
}]),
check_like=True)
rd = dbi.render_list()
assert sorted(rd.keys()) == sorted(['internal_id', 'document_id'])