yeeetty/AI4SE-proj1-checkpoints

T5 for Code Review (CSCI 680 Assignment 1)

This repository contains T5-small-style sequence-to-sequence checkpoints trained for two code-review tasks: editing Java code from a review comment, and generating a review comment from code. The model was initialized from a T5 configuration with random weights; it is not a fine-tune of a pretrained public T5 checkpoint. A custom tokenizer was trained for this project.

The repository contains multiple checkpoints. Choose the subfolder that matches your use case:

Subfolder	Description
finetuning/pt/best	Recommended multitask checkpoint: pretrained on the project corpus, then jointly fine-tuned on code editing and code-to-comment examples.
finetuning/pt-st/best	Pretrained, then fine-tuned only for code editing. Use only for the code-edit task.
finetuning/nopt/best	Random-initialized baseline, jointly fine-tuned on both tasks without the additional corpus pretraining stage.
pretraining/best	Checkpoint selected by held-out span-corruption loss, before supervised fine-tuning.
pretraining/latest	Final pretraining checkpoint and saved training state at step 200,000.
tokenizer	Project's custom tokenizer files, shared by the checkpoints.

Intended Use

These models are research and coursework artifacts for experimentation with code review. They are not general-purpose code assistants and should not be treated as authoritative reviewers. In particular, generated edits may be incorrect, incomplete, or fail to compile; generated comments may be inaccurate or miss important issues. Review and test all generated code before use. The models are not validated for security-critical, safety-critical, or production decisions.

The supervised models were evaluated on Java code-review tasks. Their intended inputs follow the assignment dataset's source format, preceded by one of these task prefixes:

• code&comment2code: for code editing from code and a reviewer comment.
• code2comment: for generating a review comment from code.

Use the source text and formatting expected by the task data. This is not a chat model and does not use a conversational chat template.

Architecture and Training

• Architecture: encoder-decoder T5ForConditionalGeneration, with 6 encoder layers, 6 decoder layers, 8 attention heads, d_model=512, d_ff=2048, and a 32,000-token vocabulary (approximately 60M parameters).
• Tokenizer: project-trained 32k ByteLevel BPE. It includes T5 sentinel tokens and the task/code markers <START>, <END>, <code>, </code>, <technical_language>, and </technical_language>.
• Corpus pretraining: span corruption on the project's curated code-and-comment corpus for 200,000 optimizer steps, using noise density 0.15 and mean span length 3.
• Supervised fine-tuning: the joint runs used a deterministic 50/50 task mixture, maximum 50,000 steps, batch size 2 with gradient accumulation 16 (effective batch size 32), learning rate 1e-4, seed 680, and validation-based early stopping. The single-task ablation used the same fine-tuning settings but trained only on code editing.
• Decoding used for reported test results: greedy generation (num_beams=1, sampling disabled), up to 512 new tokens.

Evaluation

The following are the project's test-set results, with 5,000 examples per task. Exact match normalizes whitespace. Code-edit similarity is CodeBLEU; code-to-comment similarity is BLEU-4. Similarity scores do not establish semantic correctness.

Task	Checkpoint	Exact match	Similarity
Code edit	finetuning/nopt/best	30/5,000 (0.60%)	CodeBLEU 0.5393
Code edit	finetuning/pt/best	197/5,000 (3.94%)	CodeBLEU 0.8042
Code edit	finetuning/pt-st/best	300/5,000 (6.00%)	CodeBLEU 0.8049
Code to comment	finetuning/nopt/best	0/5,000 (0.00%)	BLEU-4 0.3226
Code to comment	finetuning/pt/best	0/5,000 (0.00%)	BLEU-4 0.5622

Pretraining improved similarity scores on both tasks in the joint-model comparison, but did not produce exact-match code-to-comment predictions on this test set. The code-edit-only model performed slightly better than the joint model on the code-edit test metrics; because its task mixture differs, this ablation does not isolate cross-task interference. Results are specific to this assignment's dataset, training setup, and evaluation procedure; they are not a controlled reproduction of another paper's results.

Usage

Install compatible versions of transformers and torch, then load a fine-tuned checkpoint and the tokenizer from their repository subfolders:

import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer

repo_id = "yeeetty/AI4SE-proj1-checkpoints"
checkpoint = "finetuning/pt/best"

tokenizer = AutoTokenizer.from_pretrained(repo_id, subfolder="tokenizer")
model = AutoModelForSeq2SeqLM.from_pretrained(repo_id, subfolder=checkpoint)
model.eval()

# Paste the source field formatted as in the assignment dataset.
source_text = input("Task source: ")
source = "code&comment2code: " + source_text
inputs = tokenizer(
    source,
    return_tensors="pt",
    truncation=True,
    max_length=512,
)

with torch.inference_mode():
    output_ids = model.generate(
        **inputs,
        do_sample=False,
        num_beams=1,
        max_new_tokens=512,
    )

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

For code-to-comment generation, change the prefix to code2comment: . To load the code-edit-only ablation, set checkpoint to finetuning/pt-st/best. The nopt checkpoint is available at finetuning/nopt/best.

Data and Limitations

Pretraining used a curated corpus mined from public GitHub repositories selected with the project's GitHub Search query and filtering procedure. The supervised fine-tuning and evaluation use the course-provided code-review dataset. The project reports 448,533 instances in its final curated corpus. See the corpus repository and its data card for provenance, filtering, deduplication, decontamination, and known limitations.

The model has not been evaluated across programming languages, repositories, or review cultures beyond the stated assignment setup. Exact match is a strict metric, while BLEU and CodeBLEU measure reference similarity rather than correctness.