---
library_name: peft
license: bigcode-openrail-m
base_model: bigcode/starcoderbase-1b
tags:
- base_model:adapter:bigcode/starcoderbase-1b
- lora
- transformers
pipeline_tag: text-generation
model-index:
- name: peft-starcoder-lora-apple
  results: []
datasets:
- smangrul/hf-stack-v1
language:
- en
---

# Code Completion with StarCoder

A Python implementation of an AI-powered code completion system using the StarCoder base model with LoRA fine-tuning. This project provides a lightweight yet powerful code completion capability that can be trained on custom datasets.

##  Features

- **Fill-in-the-Middle (FIM) Capability**: Handles both prefix-suffix code completion and middle-context completion
- **LoRA Fine-tuning**: Efficient parameter-efficient fine-tuning using Low-Rank Adaptation
- **Modular Architecture**: Clean separation between settings, model components, and training logic
- **Customizable Training**: Easily adjust hyperparameters through the settings file
- **Apple Silicon Support**: Optimized for running on Apple MPS devices / Note-Since currenlty huggingface does'nt support MLX backend out of the box, hence training runs on mlx backend is slow when compare to cuda device.

##  Requirements

- Python 3.8+
- PyTorch 2.0+
- Transformers 4.30+
- PEFT (Parameter-Efficient Fine-Tuning)
- Datasets
- Accelerate
- BitsAndBytes (for quantization)

## 🛠 Installation

```bash
# Clone the repository
git clone https://github.com/deep-learner-ConfigurableAI/code-completion.git
cd code-completion

# Create a virtual environment
python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

# Install dependencies
pip install torch transformers datasets accelerate bitsandbytes peft tqdm
```

##  Project Structure

```
code-completion/
├── LICENSE
├── README.md
├── src/
│   ├── main.py           # Entry point for the application
│   ├── settings.py       # Configuration settings
│   ├── model.py          # Core model implementation
│   └── runner.py         # Training and inference logic
```

##  Configuration

All model and training configurations are centralized in `src/settings.py`. Key parameters include:

- Model checkpoint (`MODEL`)
- Training dataset (`DATASET`)
- Sequence length (`SEQ_LENGTH`)
- Training parameters (batch size, learning rate, etc.)
- LoRA configuration (rank, alpha, target modules)
- FIM transformation settings

##  Usage

### Training

To train the model on your dataset:

1. Update the `settings.py` file with your desired configuration
2. Uncomment the `train_model()` line in `main.py`
3. Run the following command:

```bash
cd src
python main.py
```

Snapshot of Trainig Run:
TrainOutput(global_step=1000, training_loss=0.7857105331420898, metrics={'train_runtime': 626.5932, 'train_samples_per_second': 12.767, 'train_steps_per_second': 1.596, 'train_tokens_per_second': 25841.328, 'total_flos': 9.961198190592e+16, 'train_loss': 0.7857105331420898, 'epoch': 0.8176614881439084})


### Inference

To use the model for code completion:

1. Ensure you have a trained model or use the provided checkpoint
2. Uncomment the `code_completion_demo()` line in `main.py`
3. Run:

```bash
cd src
python main.py
```

### Custom Inference

You can also use the model programmatically:

```python
from model import load_model_tokenizer, get_code_completion
from runner import load_model_for_inference

# Load model
model, tokenizer = load_model_for_inference()

# Example code completion
prefix = "def calculate_total(items):"
suffix = "    return total"
completed_code = get_code_completion(model, tokenizer, prefix, suffix)
print(completed_code)
```

##  How It Works

1. **Fill-in-the-Middle (FIM)**: The model is trained to predict missing code in the middle of two context pieces (prefix and suffix).

2. **LoRA Fine-tuning**: Instead of fine-tuning all parameters, we use LoRA to efficiently adapt the pre-trained StarCoder model.

3. **Dataset Processing**: The training process formats the dataset into fixed-length chunks with FIM transformations applied.

4. **Constant Length Dataset**: For efficient training, we process examples into a constant length format.

##  Performance

The model's performance depends on:
- The quality and size of the training dataset
- The hyperparameters used (especially LoRA rank and learning rate)
- The number of training steps

##  Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

##  License

This project is licensed under the MIT License - see the LICENSE file for details.

##  Acknowledgments

- [BigCode Project](https://www.bigcode-project.org/) for the StarCoder base model
- [Hugging Face](https://huggingface.co/) for their excellent Transformers library
- [PEFT](https://github.com/huggingface/peft) for the efficient fine-tuning implementation