ShimizuYuki/Marvel_network

1import csv2import os3import datasets4 5_DESCRIPTION = """6This is a dataset for Marvel universe social network, which contains the relationships between Marvel heroes.7"""8_CITATION = """\9@article{alberich2002marvel,10  title={Marvel Universe looks almost like a real social network},11  author={Alberich, Ricardo and Miro-Julia, Joe and Rossell{\'o}, Francesc},12  journal={arXiv preprint cond-mat/0202174},13  year={2002}14}15"""16 17_HOMEPAGE = "https://huggingface.co/datasets/ShimizuYuki/Marvel_network"18 19_LICENSE = "afl-3.0"20 21_URLS = {22    "adjacency_list": "https://drive.google.com/uc?id=1wcINfLn25tMIVJcp6MtxSNR7QNF8GI_D",23    "hero_hero_comic": "https://drive.google.com/uc?id=1wel0zjoa8GvBo255dlX7cVOPF9XbvQrI",24}25 26class Marvel(datasets.GeneratorBasedBuilder):27 28    VERSION = datasets.Version("1.0.1")29 30    BUILDER_CONFIGS = [31        datasets.BuilderConfig(name="adjacency_list", version=VERSION, description="This is a adjacency list for this network"),32        datasets.BuilderConfig(name="hero_hero_comic", version=VERSION, description="This adds comic imformation to adjacency list"),33    ]34 35    DEFAULT_CONFIG_NAME = "adjacency_list"36 37    def _info(self):38        if self.config.name == "adjacency_list":  # This is the name of the configuration selected in BUILDER_CONFIGS above39            features = datasets.Features(40                {41                    "hero1": datasets.Value("string"),42                    "hero2": datasets.Value("string"),43                    "counts": datasets.Value("int64")44                    # These are the features of your dataset like images, labels ...45                }46            )47        else:  # This is an example to show how to have different features for "first_domain" and "second_domain"48            features = datasets.Features(49                {50                    "hero1": datasets.Value("string"),51                    "hero2": datasets.Value("string"),52                    "comic": datasets.Value("string")53                    # These are the features of your dataset like images, labels ...54                }55            )56        return datasets.DatasetInfo(57            description=_DESCRIPTION,58            features=features,59            homepage=_HOMEPAGE,60            license=_LICENSE,61            citation=_CITATION,62        )63 64 65    def _split_generators(self, dl_manager):66        urls = _URLS[self.config.name]67        data_file = dl_manager.download(urls)68        return [69            datasets.SplitGenerator(70                name = "train",71                gen_kwargs = {72                    "filepath": data_file,73                },74            )75        ]76 77    def _generate_examples(self, filepath):78        """Generates examples as dictionaries."""79        with open(filepath, encoding="utf-8") as csv_file:80            reader = csv.DictReader(csv_file)81            for id_, row in enumerate(reader):82                if self.config.name == "adjacency_list":83                    yield id_, {84                        "hero1": row["hero1"],85                        "hero2": row["hero2"],86                        "counts": int(row["counts"]),87                    }88                else:89                    yield id_, {90                        "hero1": row["hero1"],91                        "hero2": row["hero2"],92                        "comic": row["comic"],93                    }