malteos/paperswithcode-aspects

1from __future__ import absolute_import, division, print_function2 3import json4import os5import sys6 7import datasets8from pyarrow import csv9 10_DESCRIPTION = """Papers with aspects from paperswithcode.com dataset"""11 12_HOMEPAGE = "https://github.com/malteos/aspect-document-embeddings"13 14_CITATION = '''@InProceedings{Ostendorff2022,15  title = {Specialized Document Embeddings for Aspect-based Similarity of Research Papers},16  booktitle = {Proceedings of the {ACM}/{IEEE} {Joint} {Conference} on {Digital} {Libraries} ({JCDL})},17  author = {Ostendorff, Malte and Blume, Till, Ruas, Terry and Gipp, Bela and Rehm, Georg},18  year = {2022},19}'''20 21DATA_URL = "http://datasets.fiq.de/paperswithcode_aspects.tar.gz"22 23DOC_A_COL = "from_paper_id"24DOC_B_COL = "to_paper_id"25LABEL_COL = "label"26 27# binary classification (y=similar, n=dissimilar)28LABEL_CLASSES = labels = ['y', 'n']29 30ASPECTS = ['task', 'method', 'dataset']31 32 33def get_train_split(aspect, k):34    return datasets.Split(f'fold_{aspect}_{k}_train')35 36 37def get_test_split(aspect, k):38    return datasets.Split(f'fold_{aspect}_{k}_test')39 40 41class PWCConfig(datasets.BuilderConfig):42    def __init__(self, features, data_url, aspects, **kwargs):43        super().__init__(version=datasets.Version("0.1.0"), **kwargs)44        self.features = features45        self.data_url = data_url46        self.aspects = aspects47 48 49class PWCAspects(datasets.GeneratorBasedBuilder):50    """Paper aspects dataset."""51 52    BUILDER_CONFIGS = [53        PWCConfig(54            name="docs",55            description="document text and meta data",56            # Metadata format from paperswithcode.com57            # see https://github.com/paperswithcode/paperswithcode-data58            features={59                "paper_id": datasets.Value("string"),60                "paper_url": datasets.Value("string"),61                "title": datasets.Value("string"),62                "abstract": datasets.Value("string"),63                "arxiv_id": datasets.Value("string"),64                "url_abs": datasets.Value("string"),65                "url_pdf": datasets.Value("string"),66                "aspect_tasks": datasets.Sequence(datasets.Value('string', id='task')),67                "aspect_methods": datasets.Sequence(datasets.Value('string', id='method')),68                "aspect_datasets": datasets.Sequence(datasets.Value('string', id='dataset')),69            },70            data_url=DATA_URL,71            aspects=ASPECTS,72        ),73        PWCConfig(74            name="relations",75            description=" relation data",76            features={77                DOC_A_COL: datasets.Value("string"),78                DOC_B_COL: datasets.Value("string"),79                LABEL_COL: datasets.Value("string"),80            },81            data_url=DATA_URL,82            aspects=ASPECTS,83        ),84    ]85 86    def _info(self):87        return datasets.DatasetInfo(88            description=_DESCRIPTION + self.config.description,89            features=datasets.Features(self.config.features),90            homepage=_HOMEPAGE,91            citation=_CITATION,92        )93 94    def _split_generators(self, dl_manager):95        arch_path = dl_manager.download_and_extract(self.config.data_url)96 97        if "relations" in self.config.name:98            train_file = "train.csv"99            test_file = "test.csv"100 101            generators = []102 103            # for k in [1, 2, 3, 4]:104            for aspect in self.config.aspects:105                for k in ["sample"] + [1, 2, 3, 4]:106                    folds_path = os.path.join(arch_path, 'folds', aspect, str(k))107                    generators += [108                        datasets.SplitGenerator(109                            name=get_train_split(aspect, k),110                            gen_kwargs={'filepath': os.path.join(folds_path, train_file)}111                        ),112                        datasets.SplitGenerator(113                            name=get_test_split(aspect, k),114                            gen_kwargs={'filepath': os.path.join(folds_path, test_file)}115                        )116                    ]117            return generators118 119        elif "docs" in self.config.name:120            # docs121            docs_file = os.path.join(arch_path, "docs.jsonl")122 123            return [124                datasets.SplitGenerator(name=datasets.Split('docs'), gen_kwargs={"filepath": docs_file}),125            ]126        else:127            raise ValueError()128 129    @staticmethod130    def get_dict_value(d, key, default=None):131        if key in d:132            return d[key]133        else:134            return default135 136    def _generate_examples(self, filepath):137        """Generate docs + rel examples."""138 139        if "relations" in self.config.name:140            df = csv.read_csv(filepath).to_pandas()141 142            for idx, row in df.iterrows():143                yield idx, {144                    DOC_A_COL: str(row[DOC_A_COL]),145                    DOC_B_COL: str(row[DOC_B_COL]),146                    LABEL_COL: row['label'],  # !!! labels != label147                }148 149        elif self.config.name == "docs":150            with open(filepath, 'r') as f:151                for i, line in enumerate(f):152                    doc = json.loads(line)153                    # extract feature keys from doc154                    features = {k: doc[k] if k in doc else None for k in self.config.features.keys()}155 156                    yield i, features157