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

This is an enriched version of the Code4ML dataset, a large-scale corpus of annotated Python code snippets, competition summaries, and data descriptions sourced from Kaggle. The initial release includes approximately 2.5 million snippets of machine learning code extracted from around 100,000 Jupyter notebooks. A portion of these snippets has been manually annotated by human assessors through a custom-built, user-friendly interface designed for this task. The original dataset is organized into multiple CSV files, each containing structured data on different entities: code_blocks.csv: Contains raw code snippets extracted from Kaggle. kernels_meta.csv: Metadata for the notebooks (kernels) from which the code snippets were derived. competitions_meta.csv: Metadata describing Kaggle competitions, including information about tasks and data. markup_data.csv: Annotated code blocks with semantic types, allowing deeper analysis of code structure. vertices.csv: A mapping from numeric IDs to semantic types and subclasses, used to interpret annotated code blocks. Table 1. code_blocks.csv structure Column Description code_blocks_index Global index linking code blocks to markup_data.csv. kernel_id Identifier for the Kaggle Jupyter notebook from which the code block was extracted. code_block_id Position of the code block within the notebook. code_block The actual machine learning code snippet. Table 2. kernels_meta.csv structure Column Description kernel_id Identifier for the Kaggle Jupyter notebook. kaggle_score Performance metric of the notebook. kaggle_comments Number of comments on the notebook. kaggle_upvotes Number of upvotes the notebook received. kernel_link URL to the notebook. comp_name Name of the associated Kaggle competition. Table 3. competitions_meta.csv structure Column Description comp_name Name of the Kaggle competition. description Overview of the competition task. data_type Type of data used in the competition. comp_type Classification of the competition. subtitle Short description of the task. EvaluationAlgorithmAbbreviation Metric used for assessing competition submissions. data_sources Links to datasets used. metric type Class label for the assessment metric. Table 4. markup_data.csv structure Column Description code_block Machine learning code block. too_long Flag indicating whether the block spans multiple semantic types. marks Confidence level of the annotation. graph_vertex_id ID of the semantic type. The dataset allows mapping between these tables. For example: code_blocks.csv can be linked to kernels_meta.csv via the kernel_id column. kernels_meta.csv is connected to competitions_meta.csv through comp_name. To maintain quality, kernels_meta.csv includes only notebooks with available Kaggle scores. In addition, data_with_preds.csv contains automatically classified code blocks, with a mapping back to code_blocks.csvvia the code_blocks_index column. Code4ML 2.0 Enhancements The updated Code4ML 2.0 corpus introduces kernels extracted from Meta Kaggle Code. These kernels correspond to the kaggle competitions launched since 2020. The natural descriptions of the competitions are retrieved with the aim of LLM. Notebooks in kernels_meta2.csv may not have a Kaggle score but include a leaderboard ranking (rank), providing additional context for evaluation. Applications The Code4ML 2.0 corpus is a versatile resource, enabling training and evaluation of models in areas such as: Code generation Code understanding Natural language processing of code-related tasks