CreativeInstruct

This is the code for CreativeInstruct from CREATIVEINSTRUCT: Scalably Teaching LLMs to Balance Quality, Creativity, and Diversity (Arxiv) (link to come).

The repo is currently under construction, additional models and data will be uploaded soon.

Setup

1. Create a Python environment and install dependencies:

   pip install -r requirements.txt

2. We use BACo for data generation and evaluation. Clone the repo and either set up the environment specified there, or reuse the environment from this repository.

Data Generation

Follow the steps documented in BACo to generate data.

To generate with our data (or your own dataset), first add your dataset to BACo's home directory as a folder containing a data.json file with your prompts. For example:

base-align-collab/
└── tulu/
    └── data.json

Then edit base-align-collab/src/baco/data/dataset_utils.py to register the dataset:

1. Add a prompt template to the PROMPTS dict:

"tulu": {
    "system": SYSTEM_PROMPT_INSTRUCT,
    "system_aligned": SYSTEM_PROMPT_INSTRUCT,
    "question": "Instruction: ",
    "answer_start": ""
}

2. Add a dataset loader function:

def get_dataset_tulu(num_sample=None, **kwargs):
    with open("tulu/data.json", "r", encoding="utf-8") as fd:
        narratives = json.load(fd)
    random.seed(RANDOM_SEED)
    all_prompts = [item["instruction"] for item in narratives]
    if num_sample is not None and num_sample < len(all_prompts):
        prompts = random.sample(all_prompts, num_sample)
    else:
        prompts = all_prompts
    input_data = [{"context": "", "input": prompt} for prompt in prompts]
    output_data = ["" for _ in prompts]  # No gold output for this dataset
    return input_data, output_data, "prompt", "dummy_output"

3. Register the dataset in get_dataset():

elif dataset_name == "tulu":
    dataset = get_dataset_tulu(num_sample, **kwargs)

4. Register the dataset in parse_pred_ans():

elif dataset_name in [
    "novelty-bench-curated-debug",
    "novelty-bench-curated",
    "storyarc",
    "wildchat",
    "tulu",
]:
    return parse_pred_ans_justeval(
        questions=questions,
        ans_pred=ans_pred,
        ans_gold=ans_gold,
        raw_scores=scores,
        nudging_words=nudging_words,
        score_aspects=SCORE_ASPECTS,
        print_aggregated_metric=print_aggregated_metric,
    )

Running Generation

Once your dataset is registered, save the following as a script (e.g. generate.sh) inside ~/base-align-collab and run it:

#!/bin/bash
cd ~/base-align-collab

OUTPUT_DIR=~/gen_data
mkdir -p "$OUTPUT_DIR"

# --------------------------
# Start servers
# --------------------------
BASE_PORT=8000
ALIGNED_PORT=8123

echo "Starting base model server..."
bash scripts/serve_base_1gpu.sh --port $BASE_PORT > base_server.log 2>&1 &
BASE_PID=$!
echo "Base PID: $BASE_PID"

echo "Starting aligned model server..."
bash scripts/serve_aligned_1gpu.sh --port $ALIGNED_PORT > aligned_server.log 2>&1 &
ALIGNED_PID=$!
echo "Aligned PID: $ALIGNED_PID"

# Add a sleep or health-check loop here if your servers need
# warm-up time before accepting requests.

# --------------------------
# Run BaCo
# --------------------------
echo "Running BaCo..."
bash scripts/run_baco.sh "$OUTPUT_DIR" "$BASE_PORT" "$ALIGNED_PORT"
BACO_EXIT=$?

# --------------------------
# Cleanup
# --------------------------
echo "Shutting down servers..."
kill $BASE_PID 2>/dev/null && echo "Killed base PID $BASE_PID" || echo "Base already gone"
kill $ALIGNED_PID 2>/dev/null && echo "Killed aligned PID $ALIGNED_PID" || echo "Aligned already gone"

exit $BACO_EXIT

scripts/run_baco.sh should invoke the actual generation call. For example, to generate with Llama:

python -m src.baco.inference.run \
    --dataset_name "tulu" \
    --base_host "http://localhost:8000/v1" \
    --aligned_host "http://localhost:8123/v1" \
    --base_model "meta-llama/Meta-Llama-3-8B" \
    --aligned_model "meta-llama/Meta-Llama-3-8B-Instruct" \
    --num_sample 4000 \
    --generations_per_sample 3 \
    --max_token_total 512 \
    --base_temperature 1.0 \
    --aligned_temperature 1.0 \
    --exp baco \
    --rerun --num_threads 10 \
    --router prob+punc \
    --output_root_dir "<path_to_output_dir>" \
    --top_prob_thres 0.1

Formatting Instruction-Tuning Data

Once data generation has finished, format the output into instruction-tuning data by running ~/CreativeInstruct/creative_instruct/data_utils/create_inst_data.py. Note that the input path may differ based on the naming used in the previous step (e.g. dataset name, model names, run number).

python create_inst_data.py \
    --input "~/baco_exp/ckpts/baco/runs/4/outputs_verbose/tulu/baco_base_Llama-3-8B_align_Llama-3-8B-Instruct/prob+punc_thres_0.1/all_info/input-num_4000_samples.json" \
    --output "<path_to_output_file.json>"

Train and Generate with CreativeInstruct

To train run:

python ./creative_instruct/train/fine_tune_llama.py \
    --data_path /scratch/11147/as5957/CreativeInstruct/creative_instruct/dummy_data/dummy_train.json \ #replace with full data 
    --output_dir ./creative-lora-dummy-test \ 
    --num_epochs 1 \ #change to 10 for real test
    --batch_size 1 \ #leave at default
    --val_split 0 \  #leave at default
    --skip_test \ #leave out

To generate run:

python ./creative_instruct/generate/generate_llama.py \
    --data_path ./creative_instruct/dummy_data/test_prompt.json \
    --output_file ./creative_instruct/dummy_data/test_output.json \ 
    --base_model meta-llama/Meta-Llama-3-8B-Instruct \
    --lora_path ./creative-lora-dummy-test \
    --num_generations 1 \ #change to 10
    --max_new_tokens 150 \ #change to 600
    --batch_size 1 \ #change to 4-8 
    --do_sample

To format for evals run:

python ./creative_instruct/data_utils/format_data.py \
    --input_file ./creative_instruct/dummy_data/test_output.json \
    --save_dir ./creative_instruct/dummy_data \
    --output_name test_output_formatted.jsonl \
    --model_name creative-lora-dummy-test

Evaluate with BACo

Follow evaluation script https://github.com/YichenZW/base-align-collab#evaluation to run automatic lexical and semantic metrics

Evaluate with LLM GED, Quality Metrics and WQRM

To evaluate LLM GED, and quality metrics, run below to generate a jsonl file to upload to OpenAI batch api

python ./eval_prompts/prompts.py \
    --input_file ./creative_instruct/dummy_data/test_output_formatted.jsonl \
    --output_file ./creative_instruct/dummy_data/llm_evaluation_batch.jsonl

For WQRM scores run

python ./eval_prompts/wqrm_scorer.py ./creative_instruct/dummy_data/test_output_formatted.jsonl