Code4ML: a Large-scale Dataset of annotated Machine Learning Code

We present Code4ML: a Large-scale Dataset of annotated Machine Learning Code, a corpus of Python code snippets, competition summaries, and data descriptions from Kaggle. The corpus consists of ≈ 2.5 million snippets of ML code collected from ≈ 100 thousand Jupyter notebooks. A representative fraction of the snippets is annotated by human assessors through a user-friendly interface specially designed for that purpose. The data is organized as a set of tables in CSV format. It includes several central entities: raw code blocks collected from Kaggle (code_blocks.csv), kernels (kernels_meta.csv) and competitions meta information (competitions_meta.csv). Manually annotated code blocks are presented as a separate table (murkup_data.csv). As this table contains the numeric id of the code block semantic type, we also provide a mapping from the id to semantic class and subclass (vertices.csv). Snippets information (code_blocks.csv) can be mapped with kernels meta-data via kernel_id. Kernels metadata is linked to Kaggle competitions information through comp_name. <em>To ensure the quality of the data kernels_meta.csv includes only notebooks with an available Kaggle score.</em> <em>Automatic </em>classification of code_blocks are stored in data_with_preds.csv. The mapping of this table with code_blocks.csv can be doe through <em>code_blocks_index</em> column, which corresponds to code_blocks indices. The corpus can help solve various problems, including code synthesis from a prompt in natural language, code autocompletion, and semantic code classification.