The model takes a Python function as input (embedded in a natural-language prompt) and generates a corresponding PyTest unit test with at least one assertion.

This model was trained as part of a Bachelor’s thesis on AI-powered unit test generation.

Intended uses & limitations

Intended uses

Automatic generation of baseline PyTest unit tests
Educational and research purposes
Accelerating test writing during development

Limitations

Generated tests may require human review
Not optimized for edge cases or exception-heavy logic
Designed for single-function input, not entire files

Training and evaluation data

Dataset format: JSONL (function → test pairs)
Train split: 7,731 examples
Evaluation split: 859 examples (10%)
Random seed: 42

Training procedure

Training hyperparameters

The following hyperparameters were used during training:

learning_rate: 5e-5
lr_scheduler_type: cosine
warmup_ratio: 0.1
weight_decay: 0.01
label_smoothing_factor: 0.1
per_device_train_batch_size: 4
gradient_accumulation_steps: 4
(effective batch size = 16)
max source length: 512
max target length: 256
num_epochs: 6
device: Apple Silicon (MPS)

Training results

The best checkpoint was selected based on the lowest evaluation loss.

Epoch	Eval Loss
1	1.6184
2	1.6293
3	1.6267
4	1.6255
5	1.6253
6	1.6275

Best eval_loss: 1.618388295173645
Best checkpoint: checkpoint-483 (exported to repository root)

How to use

python

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

repo = "Oluwaferanmiii/codet5p-220m-pytest-generator"

tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModelForSeq2SeqLM.from_pretrained(repo)

prompt = (
    "Given this Python function:\n```python\n"
    "def add(a, b):\n    return a + b\n"
    "```, write a PyTest unit test function in Python with at least one assert "
    "statement to verify its behavior."
)

inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=256)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Citation

bibtex

@misc{akinwale2025codet5pytest,
  author  = {Akinwale, Oluwaferanmi Philip},
  title   = {CodeT5p-220m PyTest Generator},
  year    = {2025},
  note    = {Hugging Face model repository},
  url     = {https://huggingface.co/Oluwaferanmiii/codet5p-220m-pytest-generator}
}