Fsoft-AIC/the-vault-function
The Vault is a multilingual code-text dataset with over 40 million pairs covering 10 popular programming languages. It is the largest corpus containing parallel code-text data. By building upon The Stack, a massive raw code sample collection, the Vault offers a comprehensive and clean resource for advancing research in code understanding and generation. It provides a high-quality dataset that includes code-text pairs at multiple levels, such as class and inline-level, in addition to the function level. The Vault can serve many purposes at multiple levels.
192.9k
1import os2 3import pyarrow as pa4import pyarrow.parquet as pq5import datasets6 7 8# Meta infomation9_REPO_NAME = 'Fsoft-AIC/the-vault-function'10 11_DESCRIPTION = """The Vault is a multilingual code-text dataset with over 40 million pairs covering 10 popular programming languages. 12It is the largest corpus containing parallel code-text data. By building upon The Stack, a massive raw code sample collection, 13the Vault offers a comprehensive and clean resource for advancing research in code understanding and generation. It provides a 14high-quality dataset that includes code-text pairs at multiple levels, such as class and inline-level, in addition to the function level. 15The Vault can serve many purposes at multiple levels."""16 17_HOMEPAGE = "https://huggingface.co/Fsoft-AIC"18_LICENSE = "MIT License"19_CITATION = """20@article{manh2023vault,21 title={The Vault: A Comprehensive Multilingual Dataset for Advancing Code Understanding and Generation},22 author={Manh, Dung Nguyen and Hai, Nam Le and Dau, Anh TV and Nguyen, Anh Minh and Nghiem, Khanh and Guo, Jin and Bui, Nghi DQ},23 journal={arXiv preprint arXiv:2305.06156},24 year={2023}25}26"""27################################################################################################28 29# Config metadata30_LANG_TO_TEXT = {31 "python": "python",32 "c": "c",33 "c#": "c_sharp",34 "c++": "cpp",35 "go": "go",36 "java": "java",37 "javascript": "javascript",38 "php": "php",39 "ruby": "ruby",40 "rust": "rust",41} 42_LANG_CONFIGS = ["all"] + list(_LANG_TO_TEXT.keys())43 44_TEXT_TO_LANG = {}45for lang in _LANG_TO_TEXT:46 _TEXT_TO_LANG[_LANG_TO_TEXT[lang]] = lang47 48num_shard_split = {49 "train/small/ruby": 1,50 "train/small/c": 1,51 "train/small/c_sharp": 1,52 "train/small/cpp": 1,53 "train/small/go": 1,54 "train/small/java": 2,55 "train/small/javascript": 1,56 "train/small/php": 1,57 "train/small/python": 2,58 "train/small/rust": 1,59 60 "train/medium/c": 2,61 "train/medium/c_sharp": 3,62 "train/medium/cpp": 2,63 "train/medium/go": 4,64 "train/medium/java": 6,65 "train/medium/javascript": 2,66 "train/medium/php": 4,67 "train/medium/python": 9,68 "train/medium/ruby": 1,69 "train/medium/rust": 1,70 71 "train/full/c": 7,72 "train/full/c_sharp": 13,73 "train/full/cpp": 7,74 "train/full/go": 14,75 "train/full/java": 25,76 "train/full/javascript": 6,77 "train/full/php": 15,78 "train/full/python": 33,79 "train/full/ruby": 2,80 "train/full/rust": 3,81 82 "validation/ruby": 1,83 "validation/c": 1,84 "validation/c_sharp": 1,85 "validation/cpp": 1,86 "validation/go": 1,87 "validation/java": 1,88 "validation/javascript": 1,89 "validation/php": 1,90 "validation/python": 1,91 "validation/rust": 1,92 93 "test/ruby": 1,94 "test/c": 1,95 "test/c_sharp": 1,96 "test/cpp": 1,97 "test/go": 1,98 "test/java": 1,99 "test/javascript": 1,100 "test/php": 1,101 "test/python": 1,102 "test/rust": 1103 104}105_SPLIT_CONFIGS = ["all", "train", "train/small", "train/medium", "train/full", "validation", "test"]106 107################################################################################################108 109class TheVaultFunctionConfig(datasets.BuilderConfig):110 """BuilderConfig for The Vault dataset."""111 112 def __init__(self, *args, languages=["all"], split_set= ["all"], **kwargs):113 """BuilderConfig for the The Vault dataset.114 Args:115 split_set (:obj:`List[str]`): List of split set to load.116 languages (:obj:`List[str]`): List of languages to load.117 **kwargs: keyword arguments forwarded to super.118 """119 super().__init__(120 *args,121 name= "+".join([split.replace("/", "_") for split in split_set]) + "-" + "+".join([_LANG_TO_TEXT[lang] if lang in _LANG_TO_TEXT else lang for lang in languages]),122 **kwargs,123 )124 125 languages = set([lang.lower() for lang in languages])126 split_set = set([split.lower() for split in split_set]) 127 128 assert all([language in _LANG_CONFIGS for language in languages]), f"languages {languages} contains language not in {_LANG_CONFIGS}."129 assert all([split in _SPLIT_CONFIGS for split in split_set]), f"split_set {split_set} contains element not in {_SPLIT_CONFIGS}."130 131 if "all" in split_set:132 assert len(split_set)==1, f"Passed 'all' together with other split sets. {split_set}"133 if "train" in split_set and "train/full" in split_set:134 print("WARNING - Split set 'train' and 'train/full' are similar. Force to only train/full.")135 split_set.remove("train")136 if "train" in split_set or "train/full" in split_set:137 for split in split_set:138 if "train" in split and (split != "train" and split != "train/full"):139 raise ValueError(f"Split set 'train' (or 'train/full) already contains '{split}'. Please only include one.")140 141 if "all" in languages:142 assert len(languages)==1, f"Passed 'all' together with other languages. {languages}"143 else:144 languages = [_LANG_TO_TEXT[lang] for lang in languages] # Convert to text name145 146 self.languages = list(languages)147 self.split_set= list(split_set)148 149 150class TheVaultFunction(datasets.GeneratorBasedBuilder):151 """The Vault dataset."""152 153 VERSION = datasets.Version("1.0.0")154 155 BUILDER_CONFIG_CLASS = TheVaultFunctionConfig156 BUILDER_CONFIGS = [TheVaultFunctionConfig(languages=[lang], split_set=[spl]) for lang in _LANG_CONFIGS for spl in _SPLIT_CONFIGS]157 DEFAULT_CONFIG_NAME = "all-all"158 159 160 def _info(self):161 return datasets.DatasetInfo(162 description=_DESCRIPTION,163 features=datasets.Features({164 "hexsha": datasets.Value("string"),165 "repo": datasets.Value("string"),166 "path": datasets.Value("string"), 167 "license": datasets.Sequence(datasets.Value("string")),168 "language": datasets.Value("string"),169 "identifier": datasets.Value("string"),170 "return_type": datasets.Value("string"),171 "original_string": datasets.Value("string"),172 "original_docstring": datasets.Value("string"),173 "docstring": datasets.Value("string"),174 "docstring_tokens": datasets.Sequence(datasets.Value("string")),175 "code": datasets.Value("string"),176 "code_tokens": datasets.Sequence(datasets.Value("string")),177 "short_docstring": datasets.Value("string"),178 "short_docstring_tokens": datasets.Sequence(datasets.Value("string")),179 "comment": datasets.Sequence(datasets.Value("string")),180 "parameters": [181 {182 "param": datasets.Value("string"),183 "type": datasets.Value("string"),184 }185 ],186 "docstring_params":187 {188 "returns": [189 {190 "docstring": datasets.Value("string"),191 "docstring_tokens": datasets.Sequence(datasets.Value("string")),192 "type": datasets.Value("string")193 }194 ],195 "raises": [196 {197 "docstring": datasets.Value("string"),198 "docstring_tokens": datasets.Sequence(datasets.Value("string")),199 "type": datasets.Value("string")200 }201 ],202 "params": [203 {204 "identifier": datasets.Value("string"),205 "type": datasets.Value("string"),206 "docstring": datasets.Value("string"),207 "docstring_tokens": datasets.Sequence(datasets.Value("string")),208 "default": datasets.Value("string"),209 "is_optional": datasets.Value("bool")210 }211 ],212 "outlier_params": [213 {214 "identifier": datasets.Value("string"),215 "type": datasets.Value("string"),216 "docstring": datasets.Value("string"),217 "docstring_tokens": datasets.Sequence(datasets.Value("string")),218 "default": datasets.Value("string"),219 "is_optional": datasets.Value("bool")220 }221 ],222 "others": [223 {224 "identifier": datasets.Value("string"),225 "docstring": datasets.Value("string"),226 "docstring_tokens": datasets.Sequence(datasets.Value("string"))227 }228 ]229 },230 }),231 supervised_keys=None,232 homepage=_HOMEPAGE,233 license=_LICENSE,234 citation=_CITATION,235 236 )237 238 def _split_generators(self, dl_manager):239 generators = []240 split_set = self.config.split_set241 languages = self.config.languages242 243 if "all" in split_set:244 split_set = ["train/full", "validation", "test"]245 246 if "train" in split_set:247 split_set.remove('train')248 split_set = ["train/full"] + split_set249 250 if "all" in languages:251 languages = list(_LANG_TO_TEXT.values())252 253 # train_split_files = []254 for split in split_set:255 split_files = []256 for language in languages:257 num_shards = num_shard_split[f"{split}/{language}"]258 data_files = [259 f"data/{split}/{language}-{_index:05d}-of-{num_shards:05d}.parquet"260 for _index in range(num_shards)261 ]262 files = dl_manager.download(data_files)263 split_files.extend(files)264 265 # if load_full_train and "train" in split:266 # train_split_files.extend(split_files)267 # else:268 269 generators.append(270 datasets.SplitGenerator(271 name="train" if split == "train/full" else split.replace("/", "_"),272 gen_kwargs={273 "files": split_files,274 },275 ),276 )277 278 # if load_full_train and train_split_files:279 # generators = [datasets.SplitGenerator(name="train", gen_kwargs={"files": train_split_files})] + generators280 281 282 return generators283 284 def _generate_examples(self, files):285 key = 0286 for file_idx, file in enumerate(files):287 with open(file, "rb") as f:288 parquet_file = pq.ParquetFile(f)289 for batch_idx, record_batch in enumerate(parquet_file.iter_batches(batch_size=10_000)):290 pa_table = pa.Table.from_batches([record_batch])291 for row_index in range(pa_table.num_rows):292 row = pa_table.slice(row_index, 1).to_pydict()293 294 yield key, {295 "hexsha": row['hexsha'][0],296 "repo": row['repo'][0],297 "path": row['path'][0], 298 "license": row['license'][0], 299 "language": row['language'][0],300 "identifier": row['identifier'][0],301 "return_type": row['return_type'][0],302 "original_string": row['original_string'][0],303 "original_docstring": row['original_docstring'][0],304 "docstring": row['docstring'][0],305 "docstring_tokens": row['docstring_tokens'][0],306 "code": row['code'][0],307 "code_tokens": row['code_tokens'][0],308 "short_docstring": row['short_docstring'][0],309 "short_docstring_tokens": row['short_docstring_tokens'][0],310 "comment": row['comment'][0],311 "parameters": row['parameters'][0],312 "docstring_params": row['docstring_params'][0],313 } 314 key += 1