LLaMA-Factory/src/utils/pairwise.py
hiyouga 54b8ce7b63 Initial commit
Former-commit-id: 769c6ab56be0c9d26e9289f61ac54a4068d935c1
2023-05-28 18:09:04 +08:00

52 lines
1.8 KiB
Python

import torch
from typing import Dict, Sequence, Union
from .data_collator import DataCollatorForLLaMA
from .peft_trainer import PeftTrainer
from .other import get_logger
logger = get_logger(__name__)
class PairwiseDataCollatorForLLaMA(DataCollatorForLLaMA):
r"""
Data collator for pairwise data.
"""
def __call__(self, features: Sequence[Dict[str, Union[torch.Tensor, Sequence[int]]]]) -> Dict[str, torch.Tensor]:
r"""
Pads batched data to the longest sequence in the batch.
We generate 2 * n examples where the first n examples represent chosen examples and
the last n examples represent rejected examples.
"""
features = [{"input_ids": feature[key]} for key in ("accept_ids", "reject_ids") for feature in features]
return super().__call__(features)
class PairwiseTrainerForLLaMA(PeftTrainer):
r"""
Inherits PeftTrainer to compute pairwise loss.
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.can_return_loss = True # override property to return eval_loss
def compute_loss(self, model, inputs, return_outputs=False):
r"""
Computes pairwise loss. The first n examples are chosen and the last n examples are rejected.
We use score on the EOS token to represent reward of the whole sentence.
Subclass and override to inject custom behavior. It should not be directly used by external scripts.
"""
batch_size = inputs["input_ids"].size(0) // 2
_, _, values = model(**inputs)
r_accept, r_reject = values[:, -1].split(batch_size, dim=0)
loss = -torch.log(torch.sigmoid(r_accept - r_reject)).mean()
outputs = {"r_accept": r_accept, "r_reject": r_reject}
return (loss, outputs) if return_outputs else loss