codeparrot/github-code-clean

1# coding=utf-82# Copyright 2020 The HuggingFace Datasets Authors and the current dataset script contributor.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""GitHub Code clean dataset."""16 17import os18 19import pyarrow as pa20import pyarrow.parquet as pq21 22import datasets23 24_REPO_NAME = "codeparrot/github-code-clean"25 26_LANG_TO_EXTENSION = {27    "Assembly": [".asm"],28    "Batchfile": [".bat", ".cmd"],29    "C": [".c", ".h"],30    "C#": [".cs"],31    "C++": [".cpp", ".hpp", ".c++", ".h++", ".cc", ".hh", ".C", ".H"],32    "CMake": [".cmake"],33    "CSS": [".css"],34    "Dockerfile": [".dockerfile", "Dockerfile"],35    "FORTRAN": ['.f90', '.f', '.f03', '.f08', '.f77', '.f95', '.for', '.fpp'],36    "GO": [".go"],37    "Haskell": [".hs"],38    "HTML":[".html"],39    "Java": [".java"],40    "JavaScript": [".js"],41    "Julia": [".jl"],42    "Lua": [".lua"],43    "Makefile": ["Makefile"],44    "Markdown": [".md", ".markdown"],45    "PHP": [".php", ".php3", ".php4", ".php5", ".phps", ".phpt"],46    "Perl": [".pl", ".pm", ".pod", ".perl"],47    "PowerShell": ['.ps1', '.psd1', '.psm1'],48    "Python": [".py"],49    "Ruby": [".rb"],50    "Rust": [".rs"],51    "SQL": [".sql"],52    "Scala": [".scala"],53    "Shell": [".sh", ".bash", ".command", ".zsh"],54    "TypeScript": [".ts", ".tsx"],55    "TeX": [".tex"],56    "Visual Basic": [".vb"]57}58 59_LICENSES = ['mit',60 'apache-2.0',61 'gpl-3.0',62 'gpl-2.0',63 'bsd-3-clause',64 'agpl-3.0',65 'lgpl-3.0',66 'lgpl-2.1',67 'bsd-2-clause',68 'cc0-1.0',69 'epl-1.0',70 'mpl-2.0',71 'unlicense',72 'isc',73 'artistic-2.0']74 75_DESCRIPTION = """\76The GitHub Code clean dataset in a more filtered version of codeparrot/github-code dataset, it consists of 115M code files from GitHub in 32 programming \77languages with 60 extensions totaling in almost 1TB of text data. 78"""79 80_HOMEPAGE = "https://cloud.google.com/blog/topics/public-datasets/github-on-bigquery-analyze-all-the-open-source-code/"81 82 83_EXTENSION_TO_LANG = {}84for lang in _LANG_TO_EXTENSION:85    for extension in _LANG_TO_EXTENSION[lang]:86        _EXTENSION_TO_LANG[extension] = lang87 88 89        90_LANG_CONFIGS = ["all"] + list(_LANG_TO_EXTENSION.keys())91_LICENSE_CONFIGS = ["all"] + _LICENSES92        93class GithubCodeConfig(datasets.BuilderConfig):94    """BuilderConfig for the GitHub Code dataset."""95 96    def __init__(self, *args, languages=["all"], licenses=["all"], **kwargs):97        """BuilderConfig for the GitHub Code dataset.98 99        Args:100            languages (:obj:`List[str]`): List of languages to load.101            licenses (:obj:`List[str]`): List of licenses to load.102            **kwargs: keyword arguments forwarded to super.103        """104        super().__init__(105            *args,106            name="+".join(languages)+"-"+"+".join(licenses),107            **kwargs,108        )109        110        languages = set(languages)111        licenses = set(licenses)112        113        assert all([language in _LANG_CONFIGS for language in languages]), f"Language not in {_LANG_CONFIGS}."114        assert all([license in _LICENSE_CONFIGS for license in licenses]), f"License not in {_LICENSE_CONFIGS}."115        116        if "all" in languages:117            assert len(languages)==1, "Passed 'all' together with other languages."118            self.filter_languages = False119        else:120            self.filter_languages = True121            122        if "all" in licenses:123            assert len(licenses)==1, "Passed 'all' together with other licenses."124            self.filter_licenses = False125        else:126            self.filter_licenses = True127        128        self.languages = set(languages)129        self.licenses = set(licenses)130 131 132        133class GithubCode(datasets.GeneratorBasedBuilder):134    """GitHub Code dataset."""135 136    VERSION = datasets.Version("1.0.0")137    138    BUILDER_CONFIG_CLASS = GithubCodeConfig139    BUILDER_CONFIGS = [GithubCodeConfig(languages=[lang], licenses=[license]) for lang in _LANG_CONFIGS140                                                                        for license in _LICENSE_CONFIGS]141    DEFAULT_CONFIG_NAME = "all-all"142    143    144    def _info(self):145        return datasets.DatasetInfo(146            description=_DESCRIPTION,147            features=datasets.Features({"code": datasets.Value("string"),148                                        "repo_name": datasets.Value("string"),149                                        "path": datasets.Value("string"), 150                                        "language": datasets.Value("string"),151                                        "license": datasets.Value("string"),152                                        "size": datasets.Value("int32")}),153            supervised_keys=None,154            homepage=_HOMEPAGE,155            license="Multiple: see the 'license' field of each sample.",156            157        )158 159    def _split_generators(self, dl_manager):160        num_shards = 880161        data_files = [162            f"data/train-{_index:05d}-of-{num_shards:05d}.parquet"163            for _index in range(num_shards)164        ]165        files = dl_manager.download(data_files)166        return [167            datasets.SplitGenerator(168                name=datasets.Split.TRAIN,169                gen_kwargs={170                    "files": files,171                },172            ),173        ]174 175    def _generate_examples(self, files):176        key = 0177        for file_idx, file in enumerate(files):178            with open(file, "rb") as f:179                parquet_file = pq.ParquetFile(f)180                for batch_idx, record_batch in enumerate(parquet_file.iter_batches(batch_size=10_000)):181                    pa_table = pa.Table.from_batches([record_batch])182                    for row_index in range(pa_table.num_rows):183                        row = pa_table.slice(row_index, 1).to_pydict()184                        185                        lang = lang_from_name(row['path'][0])186                        license = row["license"][0]187                        188                        if self.config.filter_languages and not lang in self.config.languages:189                            continue190                        if self.config.filter_licenses and not license in self.config.licenses:191                            continue192                        193                        yield key, {"code": row['code'][0],194                                    "repo_name": row['repo_name'][0],195                                    "path": row['path'][0],196                                    "license": license,197                                    "language": lang,198                                    "size": int(row['size'][0])}    199                        key += 1200 201                        202def lang_from_name(name):203    for extension in _EXTENSION_TO_LANG:204        if name.endswith(extension):205            return _EXTENSION_TO_LANG[extension]