# Self-Alignment Optimization for Language Models

This repository contains the official implementation of **Self-Alignment Optimization for Language Models** (SAO). The project provides an end-to-end training pipeline that integrates advanced fine-tuning techniques for optimizing language model responses through self-alignment.

## Setup and Installation

To set up the environment and install all dependencies, follow these steps:

1. **Clone the repository**:

    ```bash
    git clone https://anonymous.4open.science/r/SAO-ICLR2025-Submission.git
    cd SAO
    ```

2. **Create a conda environment**:

    Make sure you have Conda installed. Then, create a new environment with Python 3.8:

    ```bash
    conda create -n sao_env python=3.8
    conda activate sao_env
    ```

3. **Install dependencies**:

    Essential dependencies are listed in `requirements.txt`. Install them using:

    ```bash
    pip install -r requirements.txt
    ```

## Training Pipeline

The pipeline involves three major steps: prompt synthesis, response generation, and response scoring. Each step is implemented in a separate Python script.

### Step 1: Generate Prompts

The first step involves synthesizing prompts based on a given template.

```bash
python code/vllm_synthesize_prompt.py --model_path <model_path> --template <template> --sample_size <sample_size> --output_path <output_path>
```

- `--model_path`: Path to the model.
- `--template`: Prompt template to be used.
- `--sample_size`: Number of samples to use. Set to `-1` to use the full dataset.
- `--output_path`: Path to save the generated prompts.

### Step 2: Generate Responses

Next, generate responses for the synthesized prompts using the model.

```bash
python code/vllm_synthesize_response.py --model_id <model_id> --input_file <input_file> --output_file <output_file> --num_responses 2
```

- `--model_id`: Model identifier.
- `--input_file`: Path to the file containing generated prompts.
- `--output_file`: Path to save the synthesized responses.
- `--num_responses`: Number of responses to generate per prompt.

### Step 3: Score Responses

The final step is to score the generated responses using a specified prompt template.

```bash
python code/vllm_synthesize_score.py --model_id <model_id> --input_file_path <input_file_path> --output_dir <output_dir> --prompt_template_path <prompt_template_path>
```

- `--model_id`: Model identifier.
- `--input_file_path`: Path to the file containing the synthesized responses.
- `--output_dir`: Directory to save the output scores.
- `--prompt_template_path`: Path to the prompt template file used for scoring.

### Run the Full Pipeline

You can directly utilize the provided shell script to run the entire end-to-end pipeline:

```bash
bash vllm_sao.sh
```

This script sequentially performs the three steps described above for each combination of model and template.

## Output

The output files will be saved in the specified `output_dir`. Each model and template combination will have its own directory containing:

- `vllm_<template>_synthesis_prompt.jsonl`: Contains the synthesized prompts.
- `vllm_<template>_synthesis_response.jsonl`: Contains the generated responses.
- `vllm_<template>_synthesis_score.jsonl`: Contains the scores for the generated responses.

## Important Notes

- Ensure that the `output_dir` specified in the script exists, or the script will create it for you.
- If an output file already exists, the script will skip the corresponding process to save time.
