dataloaders.py

import math
import random
from collections import defaultdict
from copy import deepcopy

import numpy as np
import torch
import torch.distributed as dist
from torch.utils.data import Dataset, Sampler
from utils import IGNORE_INDEX, SEP_TOKEN

random.seed(42)
np.random.seed(42)
torch.manual_seed(42)
torch.cuda.manual_seed(42)
torch.cuda.manual_seed_all(42)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False


class GroupSampler(Sampler):
    def __init__(
        self,
        batch_size,
        dataset,
        drop_last: bool = False,
    ):
        self.batch_size = batch_size
        if dataset is None:
            raise ValueError("dataset must be provided.")
        self.dataset = dataset
        self.drop_last = drop_last

        # -------- build groups --------
        group_to_indices = defaultdict(list)
        for idx, sample in enumerate(dataset):
            key = str(sample["group_id"]) + str(sample.get("role", ""))
            group_to_indices[key].append(idx)

        self.keys = list(group_to_indices.keys())
        random.shuffle(self.keys)
        self.indices = sum([group_to_indices[k] for k in self.keys], [])

    def __iter__(self):
        return iter(self.indices)

    def __len__(self):
        return len(self.indices)


class TextDataset(Dataset):
    def __init__(self, data):
        self.data = data

    def __getitem__(self, index):
        item = self.data[index]
        assert item is not None, f"[ERROR] __getitem__ got None at idx {index}"
        return item

    def __len__(
        self,
    ):
        return len(self.data)


def sft_data_collactor(batch, tokenizer):
    input_ids, labels = [], []
    opposite_input_ids, opposite_labels = [], []
    for item in batch:
        if "query" in item:
            query = item["query"]
        elif "prompt" in item:
            query = item["prompt"]
        else:
            query = item["text"][0].split(SEP_TOKEN)[0]

        assert query != "", f"[ERROR] query is empty"

        if "target" in item:
            target = item["target"]
        elif "answer" in item:
            target = item["answer"]

        query_token_ids = tokenizer.encode(query, add_special_tokens=False)
        query_token_len = len(query_token_ids)

        target_token_ids = tokenizer.encode(target, add_special_tokens=False)

        input_ids.append(
            ([tokenizer.bos_token_id] if tokenizer.bos_token_id is not None else [])
            + deepcopy(query_token_ids)
            + deepcopy(target_token_ids)
            + [tokenizer.eos_token_id]
        )
        labels.append(
            [IGNORE_INDEX]
            * (query_token_len + (1 if tokenizer.bos_token_id is not None else 0))
            + deepcopy(target_token_ids)
            + [tokenizer.eos_token_id]
        )

        if "opposite_query" in item:
            # assert item["opposite_query"] != "" and item["opposite_target"] != "", f"[ERROR] opposite_query or opposite_target is empty"
            opposite_query_token_ids = tokenizer.encode(
                item["opposite_query"], add_special_tokens=False
            )
            opposite_query_token_len = len(opposite_query_token_ids)

            opposite_target_token_ids = tokenizer.encode(
                item["opposite_target"], add_special_tokens=False
            )
            opposite_input_ids.append(
                ([tokenizer.bos_token_id] if tokenizer.bos_token_id is not None else [])
                + deepcopy(opposite_query_token_ids)
                + deepcopy(opposite_target_token_ids)
                + [tokenizer.eos_token_id]
            )
            opposite_labels.append(
                [IGNORE_INDEX]
                * (
                    opposite_query_token_len
                    + (1 if tokenizer.bos_token_id is not None else 0)
                )
                + deepcopy(opposite_target_token_ids)
                + [tokenizer.eos_token_id]
            )

    outputs = batch_padding(input_ids + opposite_input_ids, tokenizer)
    label_outputs = batch_padding(
        labels + opposite_labels, tokenizer, pad_token_id=IGNORE_INDEX
    )

    outputs["labels"] = label_outputs["input_ids"]

    if len(opposite_input_ids) == 0:
        return {
            "input_ids": torch.tensor(outputs["input_ids"], dtype=torch.long),
            "labels": torch.tensor(outputs["labels"], dtype=torch.long),
            "attention_mask": torch.tensor(
                outputs["attention_mask"], dtype=torch.float
            ),
        }
    else:
        assert len(outputs["input_ids"]) == 2 * len(input_ids)
        return {
            "input_ids": torch.tensor(
                outputs["input_ids"][: len(input_ids)], dtype=torch.long
            ),
            "labels": torch.tensor(
                outputs["labels"][: len(input_ids)], dtype=torch.long
            ),
            "attention_mask": torch.tensor(
                outputs["attention_mask"][: len(input_ids)], dtype=torch.float
            ),
            "opposite_input_ids": torch.tensor(
                outputs["input_ids"][len(input_ids) :], dtype=torch.long
            ),
            "opposite_labels": torch.tensor(
                outputs["labels"][len(input_ids) :], dtype=torch.long
            ),
            "opposite_attention_mask": torch.tensor(
                outputs["attention_mask"][len(input_ids) :], dtype=torch.float
            ),
        }


def weighted_sft_data_collactor(batch, tokenizer):
    results = sft_data_collactor(batch, tokenizer)
    weights = [item.get("weight", 1.0) for item in batch]
    rewards = [item.get("reward", 1.0) for item in batch]
    opposite_rewards = [item.get("opposite_reward", 1.0) for item in batch]
    group_index = [item.get("group_id", -1) for item in batch]
    results["weights"] = torch.tensor(weights).float()
    results["rewards"] = torch.tensor(rewards).float()
    results["opposite_rewards"] = torch.tensor(opposite_rewards).float()
    results["group_index"] = torch.tensor(group_index, dtype=torch.long)
    return results


def offline_data_collactor(batch, tokenizer, grouped=False):
    results = weighted_sft_data_collactor(batch, tokenizer)
    sft_mask = [1.0 if item.get("type", "sample") == "sft" else 0.0 for item in batch]
    results["sft_mask"] = torch.tensor(sft_mask).float()
    return results


def batch_padding(
    input_ids, tokenizer, padding="longest", max_length=None, pad_token_id=None
):
    if pad_token_id is None:
        pad_token_id = (
            tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
        )

    max_length = tokenizer.model_max_length if max_length is None else max_length

    if padding == "longest":
        max_input_length = max([len(inp_ids) for inp_ids in input_ids])
        max_length = min(tokenizer.model_max_length, max_input_length)

    outputs = {"input_ids": [], "attention_mask": []}
    for inp_ids in input_ids:
        attn_mask = [1] * len(inp_ids)
        if len(inp_ids) >= max_length:
            if tokenizer.truncation_side == "left":
                inp_ids = inp_ids[-max_length:]
                attn_mask = attn_mask[-max_length:]
            else:
                inp_ids = inp_ids[:max_length]
                attn_mask = attn_mask[:max_length]
        else:
            if tokenizer.padding_side == "left":
                inp_ids = [pad_token_id] * (max_length - len(inp_ids)) + inp_ids
                attn_mask = [0] * (max_length - len(attn_mask)) + attn_mask
            else:
                inp_ids = inp_ids + [pad_token_id] * (max_length - len(inp_ids))
                attn_mask = attn_mask + [0] * (max_length - len(attn_mask))

        outputs["input_ids"].append(deepcopy(inp_ids))
        outputs["attention_mask"].append(deepcopy(attn_mask))
    return outputs


class ReplayBuffer:
    def __init__(self, batch_size=4, max_size=10000):
        self.max_size = max_size
        self.size = 0
        self.agent_names = None
        self.opposite_agent_names = None
        self.observations = None
        self.opposite_observations = None
        self.rewards = None
        self.opposite_rewards = None
        self.next_observations = None
        self.dones = None
        self.batch_size = batch_size
        self.actions = None
        self.opposite_actions = None
        self.mc_returns = None
        self.opposite_mc_returns = None

    def sample(self, batch_size=None):
        if batch_size is None:
            batch_size = self.batch_size
        indices = np.random.choice(
            min(self.size, self.max_size), batch_size, replace=False
        )
        return {
            "agent_name": self.agent_names[indices],
            "observation": self.observations[indices],
            "opposite_observation": self.opposite_observations[indices],
            "action": self.actions[indices],
            "opposite_action": self.opposite_actions[indices],
            "reward": self.rewards[indices],
            "opposite_reward": self.opposite_rewards[indices],
            "next_observation": self.next_observations[indices],
            "done": self.dones[indices],
            "mc_return": self.mc_returns[indices],
            "opposite_mc_return": self.opposite_mc_returns[indices],
        }

    def __len__(self):
        return min(self.size, self.max_size)

    def add(
        self,
        agent_name,
        opposite_agent_name,
        observation,
        opposite_observation,
        next_observation,
        action,
        opposite_action,
        reward: np.ndarray,
        opposite_reward: np.ndarray,
        done: np.ndarray,
        mc_return=None,
        opposite_mc_return=None,
        **kwargs,
    ):
        if self.observations is None:
            self.agent_names = np.array([""] * self.max_size, dtype="object")
            self.opposite_agent_names = np.array([""] * self.max_size, dtype="object")
            self.observations = np.array([""] * self.max_size, dtype="object")
            self.opposite_observations = np.array([""] * self.max_size, dtype="object")
            self.actions = np.array([""] * self.max_size, dtype="object")
            self.opposite_actions = np.array([""] * self.max_size, dtype="object")
            self.rewards = np.empty((self.max_size, *reward.shape), dtype=reward.dtype)
            self.opposite_rewards = np.empty(
                (self.max_size, *opposite_reward.shape), dtype=opposite_reward.dtype
            )
            self.dones = np.empty((self.max_size, *done.shape), dtype=done.dtype)
            self.mc_returns = np.empty(
                (self.max_size, *mc_return.shape), dtype=mc_return.dtype
            )
            self.opposite_mc_returns = np.empty(
                (self.max_size, *opposite_mc_return.shape),
                dtype=opposite_mc_return.dtype,
            )
            self.next_observations = np.array([""] * self.max_size, dtype="object")

        assert reward.shape == (), "reward must be a scalar"
        assert opposite_reward.shape == (), "opposite_reward must be a scalar"
        assert done.shape == (), "done must be a scalar"

        insert_at = self.size % self.max_size
        self.agent_names[insert_at] = agent_name
        self.opposite_agent_names[insert_at] = opposite_agent_name
        self.observations[insert_at] = observation
        self.opposite_observations[insert_at] = opposite_observation
        self.actions[insert_at] = action
        self.opposite_actions[insert_at] = opposite_action
        self.rewards[insert_at] = reward
        self.opposite_rewards[insert_at] = opposite_reward
        self.next_observations[insert_at] = next_observation
        self.dones[insert_at] = int(done)
        self.mc_returns[insert_at] = mc_return
        self.opposite_mc_returns[insert_at] = opposite_mc_return

        self.size += 1