metadata (allenai#132)

* wip, computing metadata for datasets (so far, hash and count)
includes a command line utility to generate the metadata

allenai#66

* some metadata

* metadata

* much faster gov and gov2 processing

* base qlogs class

* metadata refactoring

* better metadata file formatting

* clueweb doc metadata

* providing *_count via metadata if not yet available
also making more *_count functions lazy

* metadata updates

* removed a global numpy import to improve package load time

* finished clirmatrix metadata

* defer to metadata to provide docstore count hints

* beta python api support

* minor metadata file formatting

* fixed integration bugs for c4 and tweets2013-ia

* added args.me metadata

Author: seanmacavaney

.github/workflows/push.yml

@@ -45,12 +45,12 @@ jobs:
       if: matrix.os == 'ubuntu-latest' || matrix.os == 'macOs-latest'
       run: |
         pip install pytest
-        pytest test/util.py test/integration/dummy.py test/integration/vaswani.py test/formats/
+        pytest test/util.py test/metadata.py test/integration/dummy.py test/integration/vaswani.py test/formats/
     - name: Test-windows with pytest
       if: matrix.os == 'windows-latest'
       shell: cmd
       run: |
         pip install pytest
-        pytest test\util.py test\integration\dummy.py test\integration\vaswani.py test\formats\
+        pytest test\util.py test\metadata.py test\integration\dummy.py test\integration\vaswani.py test\formats\
       env:
         PATH: 'C:/Program Files/zlib/bin/'

ir_datasets/__init__.py

@@ -1,3 +1,11 @@
+from enum import Enum
+class EntityType(Enum):
+    docs = "docs"
+    queries = "queries"
+    qrels = "qrels"
+    scoreddocs = "scoreddocs"
+    docpairs = "docpairs"
+    qlogs = "qlogs"
 from . import lazy_libs
 from . import log
 from . import util
@@ -10,12 +18,11 @@
 
 Dataset = datasets.base.Dataset
 
-
 def load(name):
     return registry[name]
 
 
-def _parent_id(dataset_id: str, entity_type: str) -> str:
+def parent_id(dataset_id: str, entity_type: EntityType) -> str:
     """
     Maps a dataset_id to a more general ID that shares the same entity handler (e.g., docs_handler). For example,
     for docs, "msmarco-document/trec-dl-2019/judged" -> "msmarco-document" or "wikir/en1k/test" -> "wikir/en1k".
@@ -25,44 +32,45 @@ def _parent_id(dataset_id: str, entity_type: str) -> str:
     hierarchy that has the same docs_handler instance. This function may be updated in the future to
     also use explicit links added when datasets are registered.
     """
+    entity_type = EntityType(entity_type) # validate & allow strings
     ds = load(dataset_id)
     segments = dataset_id.split("/")
-    handler = getattr(ds, f'{entity_type}_handler')()
+    handler = getattr(ds, f'{entity_type.value}_handler')()
     parent_ds_id = dataset_id
     while len(segments) > 1:
         segments.pop()
         try:
             parent_ds = load("/".join(segments))
-            if getattr(parent_ds, f'has_{entity_type}')() and getattr(parent_ds, f'{entity_type}_handler')() == handler:
+            if parent_ds.has(entity_type.value) and getattr(parent_ds, f'{entity_type.value}_handler')() == handler:
                 parent_ds_id = "/".join(segments)
         except KeyError:
             pass # this dataset doesn't exist
     return parent_ds_id
 
 
 def docs_parent_id(dataset_id: str) -> str:
-    return _parent_id(dataset_id, 'docs')
+    return parent_id(dataset_id, EntityType.docs)
 corpus_id = docs_parent_id # legacy
 
 
 def queries_parent_id(dataset_id: str) -> str:
-    return _parent_id(dataset_id, 'queries')
+    return parent_id(dataset_id, EntityType.queries)
 
 
 def qrels_parent_id(dataset_id: str) -> str:
-    return _parent_id(dataset_id, 'qrels')
+    return parent_id(dataset_id, EntityType.qrels)
 
 
 def scoreddocs_parent_id(dataset_id: str) -> str:
-    return _parent_id(dataset_id, 'scoreddocs')
+    return parent_id(dataset_id, EntityType.scoreddocs)
 
 
 def docpairs_parent_id(dataset_id: str) -> str:
-    return _parent_id(dataset_id, 'docpairs')
+    return parent_id(dataset_id, EntityType.docpairs)
 
 
 def qlogs_parent_id(dataset_id: str) -> str:
-    return _parent_id(dataset_id, 'qlogs')
+    return parent_id(dataset_id, EntityType.qlogs)
 
 
 def create_dataset(docs_tsv=None, queries_tsv=None, qrels_trec=None):

ir_datasets/commands/__init__.py

@@ -6,6 +6,7 @@
 from . import build_download_cache
 from . import build_c4_checkpoints
 from . import clean
+from . import generate_metadata
 
 COMMANDS = {
 	'doc_fifos': doc_fifos.main,
@@ -16,4 +17,5 @@
     'build_c4_checkpoints': build_c4_checkpoints.main,
     'build_download_cache': build_download_cache.main,
     'clean': clean.main,
+    'generate_metadata': generate_metadata.main,
 }

ir_datasets/commands/generate_metadata.py

@@ -0,0 +1,92 @@
+import time
+import sys
+import os
+import json
+import argparse
+from pathlib import Path
+from fnmatch import fnmatch
+import ir_datasets
+from ir_datasets.util import DownloadConfig
+
+
+_logger = ir_datasets.log.easy()
+
+
+def dataset2metadata(args):
+    dsid, data = args
+    try:
+        dataset = ir_datasets.load(dsid)
+    except KeyError:
+        return dsid, None
+    try:
+        for e in ir_datasets.EntityType:
+            if dataset.has(e):
+                if e.value not in data:
+                    parent_id = getattr(ir_datasets, f'{e.value}_parent_id')(dsid)
+                    if parent_id != dsid:
+                        data[e.value] = {'_ref': parent_id}
+                    else:
+                        with _logger.duration(f'{dsid} {e.value}'):
+                            data[e.value] = getattr(dataset, f'{e.value}_calc_metadata')()
+                    _logger.info(f'{dsid} {e.value}: {data[e.value]}')
+    except Exception as ex:
+        _logger.info(f'{dsid} {e.value} [error]: {ex}')
+        return dsid, None
+    return dsid, data
+
+
+def write_metadata_file(data, file):
+    with file.open('wt') as f:
+        # partially-formatted data; one dataset per line
+        f.write('{\n')
+        for i, key in enumerate(sorted(data.keys())):
+            if i != 0:
+                f.write(',\n')
+            f.write(f'  "{key}": {json.dumps(data[key])}')
+        f.write('\n}\n')
+
+
+def main(args):
+    parser = argparse.ArgumentParser(prog='ir_datasets generate_metadata', description='Generates metadata for the specified datasets')
+    parser.add_argument('--file', help='output file', type=Path, default=Path('ir_datasets/etc/metadata.json'))
+    parser.add_argument('--datasets', nargs='+', help='dataset IDs for which to compute metadata. If omitted, generates for all datasets present in the registry (skipping patterns)')
+
+    args = parser.parse_args(args)
+    if args.file.is_file():
+        with args.file.open('rb') as f:
+            data = json.load(f)
+    else:
+        data = {}
+
+    if args.datasets:
+        def _ds_iter():
+            for dsid in args.datasets:
+                yield dsid, data.get(dsid, {})
+        import multiprocessing
+        with multiprocessing.Pool(10) as pool:
+            for dsid, dataset_metadata in _logger.pbar(pool.imap_unordered(dataset2metadata, _ds_iter()), desc='datasets', total=len(args.datasets)):
+                if dataset_metadata is not None:
+                    data[dsid] = dataset_metadata
+        write_metadata_file(data, args.file)
+    else:
+        for dsid in ir_datasets.registry._registered:
+            dataset = ir_datasets.load(dsid)
+            brk = False
+            try:
+                _, dataset_metadata = dataset2metadata((dsid, data.get(dsid, {})))
+                if dataset_metadata is not None:
+                    data[dsid] = dataset_metadata
+            except KeyboardInterrupt:
+                _logger.info(f'KeyboardInterrupt; skipping. ctrl+c within 0.5sec to stop compute_metadata.')
+                try:
+                    time.sleep(0.5)
+                except KeyboardInterrupt:
+                    brk = True
+                    break
+            write_metadata_file(data, args.file)
+            if brk:
+                break
+
+
+if __name__ == '__main__':
+    main(sys.argv[1:])

ir_datasets/datasets/antique.py

@@ -33,7 +33,7 @@ def _init():
     documentation = YamlDocumentation('docs/antique.yaml')
     base_path = ir_datasets.util.home_path() / NAME
     dlc = DownloadConfig.context(NAME, base_path, dua=DUA)
-    collection = TsvDocs(dlc['docs'], namespace=NAME, lang='en', count_hint=403_666)
+    collection = TsvDocs(dlc['docs'], namespace=NAME, lang='en', count_hint=ir_datasets.util.count_hint(NAME))
 
     subsets = {}
     for subset in ('train', 'test'):

ir_datasets/datasets/aol_ia.py

@@ -8,7 +8,7 @@
 import ir_datasets
 from typing import NamedTuple, Tuple
 from ir_datasets.util import DownloadConfig, GzipExtract, TarExtract, finialized_file
-from ir_datasets.formats import TrecQrels, TsvQueries, DocstoreBackedDocs
+from ir_datasets.formats import TrecQrels, TsvQueries, DocstoreBackedDocs, BaseQlogs
 from ir_datasets.datasets.base import Dataset, YamlDocumentation
 
 _logger = ir_datasets.log.easy()
@@ -45,7 +45,7 @@ class AolIaDoc(NamedTuple):
     ia_url: str
 
 
-class AolQlogs:
+class AolQlogs(BaseQlogs):
     def __init__(self, dlc):
         self.dlc = dlc
 
@@ -59,9 +59,6 @@ def qlogs_iter(self):
             except EOFError:
                 pass
 
-    def qlogs_handler(self):
-        return self
-
     def qlogs_cls(self):
         return AolQlog
 
@@ -80,8 +77,9 @@ def stream(self):
         with open(self._path, 'rb') as f:
             yield f
 
-    def path(self):
-        self._manager.build()
+    def path(self, force=True):
+        if force:
+            self._manager.build()
         return self._path
 
 
@@ -101,7 +99,7 @@ def docs_store(self):
 
     def _internal_docs_store(self):
         if self._docs_store is None:
-            self._docs_store = ir_datasets.indices.PickleLz4FullStore(self._base_path/'docs.pklz4', None, AolIaDoc, 'doc_id', ['doc_id'], count_hint=1525535)
+            self._docs_store = ir_datasets.indices.PickleLz4FullStore(self._base_path/'docs.pklz4', None, AolIaDoc, 'doc_id', ['doc_id'], count_hint=ir_datasets.util.count_hint(NAME))
         return self._docs_store
 
     def _build_docs(self):

ir_datasets/datasets/aquaint.py

@@ -26,7 +26,7 @@ def _init():
     dlc = DownloadConfig.context(NAME, base_path)
     documentation = YamlDocumentation(f'docs/{NAME}.yaml')
 
-    collection = TrecDocs(dlc['docs'], encoding='utf8', path_globs=['aquaint_comp/apw/*/*.gz', 'aquaint_comp/nyt/*/*.gz', 'aquaint_comp/xie/*/*.gz'], namespace=NAME, lang='en', count_hint=1033461)
+    collection = TrecDocs(dlc['docs'], encoding='utf8', path_globs=['aquaint_comp/apw/*/*.gz', 'aquaint_comp/nyt/*/*.gz', 'aquaint_comp/xie/*/*.gz'], namespace=NAME, lang='en', count_hint=ir_datasets.util.count_hint(NAME))
 
     base = Dataset(collection, documentation('_'))