DOI: 10.5281/zenodo.15020642
title: Replication package for DRAGON: Robust Classification for Very Large Collections of Software Repositories

description: DRAGON: Multi-Label Classification

This archive contains the replication package for the DRAGON multi-label classification models, which leverage BERT-based architectures. The package includes scripts for repository mining, dataset creation, data processing, model training, and evaluation. The two main models used are DRAGON and LEGION.

Key Components:



Repository Mining: Scripts to extract repositories for dataset creation.

Dataset Preparation: Jupyter notebooks for cleaning and transforming data.

Data Processing: Conversion into a Hugging Face dataset format.

Model Training: Training scripts for DRAGON and LEGION, with configurable preprocessing options.

Evaluation: Threshold tuning and performance assessment.


Setup

Before running any commands, ensure you have the necessary dependencies installed. It is recommended to use a virtual environment:

python3 -m venv venv
source venv/bin/activate  # On Windows use `venv\Scripts\activate`
pip install -r requirements.txt


Project Structure



repository_mining/: Contains scripts for mining the initial set of repositories.



repository_mining/doc/: Includes documentation with the necessary information for repository mining.



dataset_creation/: Contains all the notebooks to be run sequentially to prepare the dataset.

multilabel_class/: Contains scripts for classification, threshold tuning, and evaluation.



multilabel_class/model_output/: trained model organized by: first dataset, then model variantion.



data/: Contains the hugginface datasets ( our dataset and LEGION dataset) ready for the training/eval.


1️⃣ Data Mining

To mine the initial set of repositories from Software Heritage, use the scripts available in the repository_mining/ folder. Detailed information and steps for repository mining can be found in:

repository_mining/doc/


2️⃣ Dataset Creation

After mining the repositories, prepare the dataset by running the Jupyter notebooks inside the dataset_creation/ folder in sequence. These notebooks handle data cleaning, transformation, and formatting necessary for model training. All the documentation needed is inside each notebook explaining every step.

3️⃣ Data Processing

Once the dataset is prepared, convert it into a Hugging Face dataset using:

python3 multilabel_class/create_dataset.py --file_path data/02_processed_datasets/2024-05-22/origin-metadata-readme_names-900000dataset_forks-cleaned.csv


4️⃣ Classification / Training

Train the DRAGON Model

After processing the dataset, train the DRAGON model with the following command:

python3 multilabel_class/tune_thresholds.py --model_type bert --model_variant focal --dataset_path data/03_huggingaceV_datasets/2024-05-22/origin-metadata-readme_names-900000dataset_forks-cleaned/dataset


Ensure Configuration is Set Correctly

Modify the configuration file multilabel_class/utils/config.py to set the following parameter to True:

DEFAULT_PREPROCESSING_PARAMS = { 
    'use_sentence_pairs': True  # If True, process as (text1, text2); if False, concatenate texts
}


Training DRAGON Without Sentence Pairs

To train DRAGON without using sentence pairs, use the same command but set use_sentence_pairs to False in the config file:

DEFAULT_PREPROCESSING_PARAMS = { 
    'use_sentence_pairs': False
}


Train DRAGON on a Benchmark Dataset

To train DRAGON on a benchmark dataset, use:

python3 multilabel_class/tune_thresholds.py --model_type bert --model_variant focal --dataset_path data/03_huggingaceV_datasets/LEGION/dataset


Ensure the use_sentence_pairs parameter is set to True in config.py.

Train LEGION on the DRAGON Dataset

To train LEGION on the DRAGON dataset, use:

python3 multilabel_class/tune_thresholds.py --model_type bert --model_variant db --dataset_path data/03_huggingaceV_datasets/2024-05-22/origin-metadata-readme_names-900000dataset_forks-cleaned/dataset


Ensure the use_sentence_pairs parameter is set to False in config.py:

DEFAULT_PREPROCESSING_PARAMS = { 
    'use_sentence_pairs': False
}


Train LEGION on a Baseline Dataset

To train LEGION on a baseline dataset, run:

python3 multilabel_class/tune_thresholds.py --model_type bert --model_variant db --dataset_path data/03_huggingaceV_datasets/LEGION/dataset


5️⃣ Model Evaluation

Once thresholds are tuned, you can evaluate the model using:

python3 multilabel_class/evaluation.py --model_type bert --model_variant focal --dataset_path data/03_huggingaceV_datasets/2024-05-22/origin-metadata-readme_names-900000dataset_forks-cleaned/dataset


This evaluation script computes standard multi-label classification metrics including:



Micro and macro F1@1..5-score

Precision@1..5 and recall@1..5


Ensure that the model variant and dataset path correspond to the previously trained model.

Recommended: Evaluation via Notebooks

We suggest an interactive and visual analysis of model performance, you can also use the provided Jupyter notebooks located in:

DRAGON_replication/multilabel_class/notebooks/


These notebooks reproduce the complete evaluation pipeline and generate additional visualizations and metrics discussed in the associated paper.

Both command-line and notebook-based evaluations ensure reproducibility and offer complementary insights into model behavior.

 

This README provides an overview of the essential steps for repository mining, dataset preparation, processing, model training, and evaluation. For further customization, refer to the configuration files and experiment with different preprocessing settings.