[Paper Review] PyTorrent: A Python Library Corpus for Large-scale Language Models

Mehdi Bahrami, N. C. Shrikanth+7 more|arXiv (Cornell University)|Oct 4, 2021

Software Engineering Research27 references4 citations

TL;DR

PyTorrent is a large-scale, publicly available corpus of 218,814 Python package libraries from PyPI and Anaconda, curated for training and fine-tuning large-scale language models in software engineering tasks. It provides source code, natural language descriptions, metadata, and a pre-trained DistilBERT-based model, enabling direct use in code search, code generation, and defect prediction with minimal infrastructure overhead.

ABSTRACT

A large scale collection of both semantic and natural language resources is essential to leverage active Software Engineering research areas such as code reuse and code comprehensibility. Existing machine learning models ingest data from Open Source repositories (like GitHub projects) and forum discussions (like Stackoverflow.com), whereas, in this showcase, we took a step backward to orchestrate a corpus titled PyTorrent that contains 218,814 Python package libraries from PyPI and Anaconda environment. This is because earlier studies have shown that much of the code is redundant and Python packages from these environments are better in quality and are well-documented. PyTorrent enables users (such as data scientists, students, etc.) to build off the shelf machine learning models directly without spending months of effort on large infrastructure. The dataset, schema and a pretrained language model is available at: https://github.com/fla-sil/PyTorrent

Motivation & Objective

• To address the lack of high-quality, curated Python library corpora for training large-scale language models in software engineering.
• To improve code reuse and comprehensibility by providing a dataset of high-quality, well-documented Python packages from PyPI and Anaconda.
• To enable researchers and practitioners to fine-tune existing models like CodeBERT, CodeXGLUE, and GraphCodeBERT without extensive data preparation or infrastructure.
• To support downstream tasks such as code retrieval, code generation, and automated program repair through a standardized, schema-compliant dataset.
• To reduce training time and computational cost by providing a pre-tokenized, truncated, and optimized dataset with a pre-trained model.

Proposed method

• The dataset was constructed by mining 218,814 Python packages from PyPI and Anaconda, extracting source code, docstrings, and metadata (e.g., package name, license, supported Python version).
• A schema compliant with CodeSearchNet was adopted to ensure compatibility with existing transformer-based models, including fields for natural language descriptions (NL), source code (PL), function names, and tokens.
• A byte-level BPE tokenizer was trained on 1 million raw Python scripts (12.35M LOC), truncated to 50 tokens per sequence to reduce computational load.
• A DistilBERT-based masked language model was trained from scratch on 1M scripts to create a pre-trained model, released as PyTorrent(v1) on Hugging Face Hub.
• The dataset includes three augmented code scenarios to support diverse downstream tasks, including code search and code generation.
• Fine-tuning experiments were conducted on CodeBERT using the PyTorrent dataset, demonstrating strong performance in multi-line code generation.

Experimental results

Research questions

• RQ1 Can a curated corpus of Python package libraries from PyPI and Anaconda improve the performance of code search and code generation models?
• RQ2 To what extent does using high-quality, well-documented packages from package managers enhance model generalization compared to mining GitHub repositories?
• RQ3 How effective is a pre-trained model trained on PyTorrent for downstream tasks like code retrieval and code generation?
• RQ4 Can the PyTorrent dataset be seamlessly integrated into existing frameworks like CodeSearchNet, CodeBERT, and GraphCodeBERT for fine-tuning?
• RQ5 What is the impact of filtering out trivial packages (≤35 LOC, cyclomatic complexity ≤10) on model training efficiency and output quality?

Key findings

• PyTorrent contains 4,058,349 files, 655 million lines of code, and 2.84 million functions, making it one of the largest curated Python code corpora to date.
• The dataset includes metadata for all 218,814 packages, such as license, supported Python versions, and package descriptions, enhancing model interpretability and usability.
• A pre-trained DistilBERT model (PyTorrent(v1)) was successfully trained from scratch on 1M scripts and released on Hugging Face Hub for immediate deployment.
• Fine-tuning CodeBERT on PyTorrent yielded strong results in multi-line code generation, demonstrating the dataset’s value for complex code synthesis tasks.
• The dataset is compatible with major code modeling frameworks, including CodeSearchNet, CodeXGLUE, and GraphCodeBERT, enabling plug-and-play integration for downstream research.
• The use of a BPE tokenizer with 56,000 tokens and sequence truncation to 50 tokens significantly reduced computational costs while preserving model performance.