Code Preprocessor

A machine learning pipeline for preprocessing and training models on source code, specifically optimized for Rust codebases.

Requirements

• Python 3.10 or higher
• CUDA-capable GPU (tested with RTX 4070)
• WSL2 or Linux environment

Installation

1. Clone the repository:

git clone https://github.com/argakiig/code-preprocessor.git
cd code-preprocessor

2. Install the package and dependencies:

pip install -e .

For development, install with additional tools:

pip install -e ".[dev]"

Usage

Basic Usage

Process and train on a Rust codebase:

code-preprocess --config config.yaml --code-path /path/to/rust/code

Configuration

Create a config.yaml file with your settings:

# Model configurationmodel:
name: "codellama/CodeLlama-7b-hf"vocab_size: 50000max_sequence_length: 256# Training configurationtraining:
batch_size: 2num_workers: 8epochs: 3gradient_accumulation_steps: 1eval_split: 0.1seed: 42# Pathspaths:
output_dir: "./output"cache_dir: "~/.cache/huggingface"# Logginglogging:
level: "INFO"file: "logs/training.log"# Weights & Biaseswandb:
project: "your-project-name"# GPU Configurationgpu:
device: "cuda"precision: "fp16"memory_efficient: true

Configuration Precedence

The tool supports both YAML configuration files and command-line arguments. When both are provided:

1. Command-line arguments take precedence over the config file settings
2. Any settings not specified in command-line arguments will fall back to the config file values
3. If a setting is not specified in either place, default values will be used

For example:

# This will use batch_size=4 from CLI, but keep other settings from config.yaml
code-preprocess --config config.yaml --code-path /path/to/code --batch-size 4

Available CLI arguments match the configuration file options and can be viewed with:

code-preprocess --help

Development Tools

The project includes several development tools:

• black: Code formatting
• isort: Import sorting
• flake8: Code linting
• mypy: Type checking
• pytest: Testing
• pre-commit: Git hooks

Set up pre-commit hooks:

pre-commit install

Run tests:

pytest

Features

• Processes Rust source code for training
• Custom tokenizer trained on Rust code
• Memory-efficient training with LoRA and 4-bit quantization
• Integrated with Weights & Biases for experiment tracking
• Automatic evaluation during training
• Best model checkpoint saving

Project Structure

code-preprocessor/
├── code_preprocessor/
│ ├── __init__.py
│ ├── __main__.py
│ ├── config.py
│ ├── processor.py
│ ├── constants.py
│ ├── models/
│ ├── parsers/
│ ├── tokenizers/
│ ├── training/
│ └── utils/
├── tests/
├── .pre-commit-config.yaml
├── pyproject.toml
└── README.md

License

MIT License