wehub-resource-sync
593b94c120
pytest / Unit Tests (push) Has been cancelled
pytest / Integration (integration_tests_a) (push) Has been cancelled
pytest / Integration (integration_tests_b) (push) Has been cancelled
pytest / Integration (integration_tests_c) (push) Has been cancelled
pytest / Integration (integration_tests_d) (push) Has been cancelled
pytest / Integration (integration_tests_e) (push) Has been cancelled
pytest / Integration (integration_tests_f) (push) Has been cancelled
pytest / Integration (integration_tests_g) (push) Has been cancelled
pytest / Integration (integration_tests_h) (push) Has been cancelled
pytest / Integration (integration_tests_i) (push) Has been cancelled
pytest / Integration (integration_tests_j) (push) Has been cancelled
pytest / Distributed (distributed_a) (push) Has been cancelled
pytest / Distributed (distributed_b) (push) Has been cancelled
pytest / Distributed (distributed_c) (push) Has been cancelled
pytest / Distributed (distributed_d) (push) Has been cancelled
pytest / Distributed (distributed_e) (push) Has been cancelled
pytest / Distributed (distributed_f) (push) Has been cancelled
pytest / Minimal Install (push) Has been cancelled
pytest / Event File (push) Has been cancelled
pytest (slow) / py-slow (push) Has been cancelled
Publish JSON Schema / publish-schema (push) Has been cancelled
97 lines
2.6 KiB
Python
97 lines
2.6 KiB
Python
"""DPO alignment training with Ludwig.
|
|
|
|
Usage:
|
|
python train_dpo.py --dataset train.csv
|
|
python train_dpo.py --dataset train.csv --epochs 3 --beta 0.05
|
|
|
|
Prerequisites:
|
|
pip install "ludwig[llm]" datasets
|
|
export HUGGING_FACE_HUB_TOKEN="<your_token>"
|
|
|
|
The dataset must have columns: prompt, chosen, rejected
|
|
Use prepare_dataset.py to produce this file from Anthropic/hh-rlhf.
|
|
"""
|
|
|
|
import argparse
|
|
import logging
|
|
import os
|
|
|
|
import yaml
|
|
|
|
from ludwig.api import LudwigModel
|
|
|
|
|
|
def build_config(epochs: int, learning_rate: float, beta: float, batch_size: int) -> dict:
|
|
raw = f"""
|
|
model_type: llm
|
|
base_model: meta-llama/Llama-3.1-8B
|
|
|
|
adapter:
|
|
type: lora
|
|
r: 16
|
|
alpha: 32
|
|
dropout: 0.05
|
|
|
|
trainer:
|
|
type: dpo
|
|
epochs: {epochs}
|
|
learning_rate: {learning_rate}
|
|
batch_size: {batch_size}
|
|
gradient_accumulation_steps: 8
|
|
beta: {beta}
|
|
|
|
input_features:
|
|
- name: prompt
|
|
type: text
|
|
|
|
output_features:
|
|
- name: chosen
|
|
type: text
|
|
|
|
backend:
|
|
type: local
|
|
"""
|
|
return yaml.safe_load(raw)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Run DPO alignment training with Ludwig.")
|
|
parser.add_argument("--dataset", required=True, help="Path to the DPO CSV (prompt, chosen, rejected).")
|
|
parser.add_argument("--epochs", type=int, default=1)
|
|
parser.add_argument("--learning_rate", type=float, default=5e-7)
|
|
parser.add_argument("--beta", type=float, default=0.1, help="KL penalty coefficient.")
|
|
parser.add_argument("--batch_size", type=int, default=2)
|
|
parser.add_argument("--experiment_name", default="hh_rlhf_dpo")
|
|
parser.add_argument("--output_dir", default="results")
|
|
args = parser.parse_args()
|
|
|
|
token = os.environ.get("HUGGING_FACE_HUB_TOKEN") or os.environ.get("HF_TOKEN")
|
|
if not token:
|
|
raise OSError(
|
|
"Set HUGGING_FACE_HUB_TOKEN (or HF_TOKEN) before running. "
|
|
"You also need access approval for meta-llama/Llama-3.1-8B."
|
|
)
|
|
|
|
config = build_config(
|
|
epochs=args.epochs,
|
|
learning_rate=args.learning_rate,
|
|
beta=args.beta,
|
|
batch_size=args.batch_size,
|
|
)
|
|
|
|
model = LudwigModel(config=config, logging_level=logging.INFO)
|
|
|
|
train_stats, preprocessed_data, output_directory = model.train(
|
|
dataset=args.dataset,
|
|
experiment_name=args.experiment_name,
|
|
output_directory=args.output_dir,
|
|
)
|
|
|
|
print(f"\nTraining complete. Results saved to: {output_directory}")
|
|
print("To upload the model to HuggingFace Hub:")
|
|
print(f" ludwig upload hf_hub -r <your_org>/<model_name> -m {output_directory}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|