Mini-SWE-RL CodeLM Final
Mini-SWE-RL CodeLM Final is a small decoder-only Python code language model trained from scratch on local ROCm hardware, then supervised fine-tuned for no-leak Python docstring generation.
This is a coursework artifact, not a production code assistant. The model is intentionally compact so it can run on an AMD RX 6600-class GPU.
Model Details
- Architecture: GPT-style decoder-only transformer
- Parameters: approximately 14M
- Vocabulary: 8,192 byte-level BPE tokens
- Context length: 512 tokens
- Layers: 6
- Attention heads: 6
- Embedding size: 384
- Dropout: 0.1
- Attention backend: causal attention (
use_sdpa=False)
- Framework: PyTorch
- Tokenizer:
tokenizers.ByteLevelBPETokenizer
The Hugging Face repository keeps the final artifacts at the repository root:
1README.md
2main.ipynb
3pretrained_code_lm.pt
4finetuned_best.pt
5finetuned_latest.pt
6tokenizer.json
7vocab.json
8merges.txt
9codexglue_nodoc_metrics_summary.csv
10codexglue_nodoc_zero_shot_predictions.csv
11codexglue_nodoc_finetuned_predictions.csv
In the local coursework workspace, the same final artifacts are consolidated under:
Intended Use
The fine-tuned checkpoint is intended for simple Python code-to-text/docstring generation experiments.
Example task:
1Write a concise English docstring for this Python function.
2
3Python function:
4<docstring-stripped Python function>
5
6Docstring:
This model should not be used for security-critical code generation, automated program repair, or production software engineering decisions.
Training Data
Pretraining
The base model was trained from scratch on Python code from codeparrot/codeparrot-clean.
The local final directory contains the tokenized pretraining bins:
1mini_swe_rl/final/data/pretrain_train.bin
2mini_swe_rl/final/data/pretrain_valid.bin
Those large pretraining bins are not uploaded to the Hugging Face model repository.
Fine-Tuning
The downstream task uses the Python split of google/code_x_glue_ct_code_to_text.
To avoid answer leakage, docstrings are removed from source code inputs using Python AST parsing before fine-tuning and evaluation. Examples are skipped if the target docstring still appears in the stripped code.
Training Procedure
Pretraining
- Objective: causal next-token prediction
- Steps: 425,000
- Batch size: 4
- Gradient accumulation: 16
- Optimizer: AdamW
- Learning rate: cosine schedule from
3e-4 to 3e-5
- Hardware target: AMD RX 6600 with ROCm
Fine-Tuning
- Objective: supervised docstring generation
- Best checkpoint step: 28,000
- Final step: 30,000
- Batch size: 8
- Gradient accumulation: 4
- Peak learning rate:
2e-5
- Minimum learning rate:
2e-6
- Best validation loss: 2.798454
Evaluation
Evaluation used 1,000 held-out CodeXGLUE Python code-to-text examples with docstring leakage removed.
| Stage | Examples | ROUGE-L | BLEU-4 | Exact Match |
|---|
| Pretrained zero-shot | 1,000 | 0.085215 | 0.005946 | 0.000 |
| Fine-tuned | 1,000 | 0.216660 | 0.039844 | 0.002 |
The fine-tuned model improves over the zero-shot pretrained baseline, but absolute scores remain modest.
Limitations
- The model is small and has a short 512-token context window.
- It often produces generic docstrings when the function behaviour requires detailed reasoning.
- Exact match is very low because generated docstrings rarely match references word-for-word.
- The checkpoint is not in Hugging Face
transformers format; it is a PyTorch checkpoint that requires the matching GPT class from main.ipynb.
- The model is not reliable for executable code generation or benchmark pass@k tasks.
Local Loading Notes
Use main.ipynb to load the tokenizer, model class, and checkpoints. In the Hugging Face repository, the best downstream checkpoint is:
The checkpoint stores:
1model state_dict
2config
3tokenizer_dir
4extra training metadata
The tokenizer files are:
1vocab.json
2merges.txt
3tokenizer.json
Citation
This model was created as part of a local small language model coursework project. Dataset citations should follow the upstream dataset cards for CodeParrot-clean and CodeXGLUE.
