LegoGPT

Installation

Prerequisites

You need a Gurobi licence to use Gurobi for stability analysis. Get a free academic licence here.

Installing as a standalone project

This repo uses the Python project manager uv. To install this repo as a standalone project:

1. Clone the repo: git clone "git@github.com:AvaLovelace1/LegoGPT.git" && cd LegoGPT.
2. Install the ImportLDraw submodule (required for rendering LEGO visualizations): git submodule init && git submodule update.
3. Some files in the ImportLDraw submodule are stored using the Git LFS system. To download these files, install Git LFS, cd into the ImportLDraw directory, and run git lfs pull.
4. Finally, install uv. A virtual environment will be created, and the remaining dependencies installed automatically, upon invoking uv run [SCRIPT_NAME].

Installing as a package

To install this repo as a package in your own Python project, run

uv add "git+ssh://git@github.com/AvaLovelace1/LegoGPT.git"

if using uv, or

pip install "git+ssh://git@github.com/AvaLovelace1/LegoGPT.git"

if using pip.

Running inference interactively

You can run inference with the fine-tuned LegoGPT model using:

uv run infer --model_name_or_path MODEL_PATH

This script starts an interactive session where you can input a prompt and get a response from the model. See uv run infer -h for a full list of options.

Fine-tuning LegoGPT

We use Hugging Face TRL with Accelerate for fine-tuning. To run fine-tuning, follow these instructions:

1. Start with a LEGO dataset with the fields "caption" and "lego". The "caption" field should contain a description of the LEGO model, and the "lego" field should contain the corresponding LEGO model, in the text format described in the paper.
2. Prepare the dataset for finetuning with uv run prepare_finetuning_dataset --input_path LEGO_DATASET_PATH --output_path FINETUNING_DATASET_PATH.
3. Download the pretrained Llama-3.2-1B-Instruct model to some directory [PRETRAINED_DIR]. IMPORTANT: Replace the config.json, special_tokens_map.json, and tokenizer_config.json files with the ones in the finetuning_config_files directory. This specifies the pad_token to be different from the eos_token, fixing a fine-tuning issue where the model will not learn to output EOS tokens properly.
4. Initialize the Accelerate config file with uv run accelerate config.
5. Run finetuning with uv run ./finetune.zsh [PRETRAINED_DIR] [OUTPUT_DIR] [RUN_NAME] [FINETUNING_DATASET_PATH]. The finetuned model will be saved to [OUTPUT_DIR]/[RUN_NAME].