macrocosm-os/code-parrot-github-code

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 dataset."""16 17import os18 19import pyarrow as pa20import pyarrow.parquet as pq21 22import datasets23 24_REPO_NAME = "codeparrot/github-code"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 dataest consists of 115M code files from GitHub in 32 programming \77languages with 60 extensions totalling in 1TB of text data. The dataset was created \78from the GitHub dataset on BiqQuery.79"""80 81_HOMEPAGE = "https://cloud.google.com/blog/topics/public-datasets/github-on-bigquery-analyze-all-the-open-source-code/"82 83 84_EXTENSION_TO_LANG = {}85for lang in _LANG_TO_EXTENSION:86    for extension in _LANG_TO_EXTENSION[lang]:87        _EXTENSION_TO_LANG[extension] = lang88 89 90        91_LANG_CONFIGS = ["all"] + list(_LANG_TO_EXTENSION.keys())92_LICENSE_CONFIGS = ["all"] + _LICENSES93        94class GithubCodeConfig(datasets.BuilderConfig):95    """BuilderConfig for the GitHub Code dataset."""96 97    def __init__(self, *args, languages=["all"], licenses=["all"], **kwargs):98        """BuilderConfig for the GitHub Code dataset.99 100        Args:101            languages (:obj:`List[str]`): List of languages to load.102            licenses (:obj:`List[str]`): List of licenses to load.103            **kwargs: keyword arguments forwarded to super.104        """105        super().__init__(106            *args,107            name="+".join(languages)+"-"+"+".join(licenses),108            **kwargs,109        )110        111        languages = set(languages)112        licenses = set(licenses)113        114        assert all([language in _LANG_CONFIGS for language in languages]), f"Language not in {_LANG_CONFIGS}."115        assert all([license in _LICENSE_CONFIGS for license in licenses]), f"License not in {_LICENSE_CONFIGS}."116        117        if "all" in languages:118            assert len(languages)==1, "Passed 'all' together with other languages."119            self.filter_languages = False120        else:121            self.filter_languages = True122            123        if "all" in licenses:124            assert len(licenses)==1, "Passed 'all' together with other licenses."125            self.filter_licenses = False126        else:127            self.filter_licenses = True128        129        self.languages = set(languages)130        self.licenses = set(licenses)131 132 133        134class GithubCode(datasets.GeneratorBasedBuilder):135    """GitHub Code dataset."""136 137    VERSION = datasets.Version("1.0.0")138    139    BUILDER_CONFIG_CLASS = GithubCodeConfig140    BUILDER_CONFIGS = [GithubCodeConfig(languages=[lang], licenses=[license]) for lang in _LANG_CONFIGS141                                                                        for license in _LICENSE_CONFIGS]142    DEFAULT_CONFIG_NAME = "all-all"143    144    145    def _info(self):146        return datasets.DatasetInfo(147            description=_DESCRIPTION,148            features=datasets.Features({"code": datasets.Value("string"),149                                        "repo_name": datasets.Value("string"),150                                        "path": datasets.Value("string"), 151                                        "language": datasets.Value("string"),152                                        "license": datasets.Value("string"),153                                        "size": datasets.Value("int32")}),154            supervised_keys=None,155            homepage=_HOMEPAGE,156            license="Multiple: see the 'license' field of each sample.",157            158        )159 160    def _split_generators(self, dl_manager):161        num_shards = 1126162        data_files = [163            f"data/train-{_index:05d}-of-{num_shards:05d}.parquet"164            for _index in range(num_shards)165        ]166        files = dl_manager.download(data_files)167        return [168            datasets.SplitGenerator(169                name=datasets.Split.TRAIN,170                gen_kwargs={171                    "files": files,172                },173            ),174        ]175 176    def _generate_examples(self, files):177        key = 0178        for file_idx, file in enumerate(files):179            with open(file, "rb") as f:180                parquet_file = pq.ParquetFile(f)181                for batch_idx, record_batch in enumerate(parquet_file.iter_batches(batch_size=10_000)):182                    pa_table = pa.Table.from_batches([record_batch])183                    for row_index in range(pa_table.num_rows):184                        row = pa_table.slice(row_index, 1).to_pydict()185                        186                        lang = lang_from_name(row['path'][0])187                        license = row["license"][0]188                        189                        if self.config.filter_languages and not lang in self.config.languages:190                            continue191                        if self.config.filter_licenses and not license in self.config.licenses:192                            continue193                        194                        yield key, {"code": row['content'][0],195                                    "repo_name": row['repo_name'][0],196                                    "path": row['path'][0],197                                    "license": license,198                                    "language": lang,199                                    "size": int(row['size'][0])}    200                        key += 1201 202                        203def lang_from_name(name):204    for extension in _EXTENSION_TO_LANG:205        if name.endswith(extension):206            return _EXTENSION_TO_LANG[extension]