torchrl 0.11.0__cp314-cp314-manylinux_2_28_aarch64.whl

torchrl/data/llm/prompt.py

@@ -0,0 +1,198 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import torch
+from tensordict import tensorclass, TensorDict
+
+from torchrl.data.llm.dataset import TensorDictTokenizer, TokenizedDatasetLoader
+
+DEFAULT_DATASET = "CarperAI/openai_summarize_tldr"
+
+
+@tensorclass
+class PromptData:
+    """A prompt dataset."""
+
+    input_ids: torch.Tensor
+    attention_mask: torch.Tensor
+    prompt_rindex: torch.Tensor
+    labels: torch.Tensor | None = None
+    logits: torch.Tensor | None = None
+    loss: torch.Tensor | None = None
+
+    def mask_label(self, pad_token_id=50256):
+        _, block_size = self.input_ids.shape
+        attention_mask = (
+            torch.arange(block_size, device=self.prompt_rindex.device)
+            < self.prompt_rindex[:, None]
+        ).to(torch.int64)
+        input_ids = torch.where(attention_mask == 1, self.input_ids, pad_token_id)
+        return self.__class__(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            prompt_rindex=self.prompt_rindex,
+            loss=self.loss,
+            batch_size=[],
+        )
+
+    @classmethod
+    def from_dataset(
+        cls,
+        split,
+        dataset_name=None,
+        max_length=550,
+        root_dir=None,
+        from_disk=False,
+        num_workers: int | None = None,
+    ):
+        """Returns a :class:`PromptData` from a dataset name.
+
+        Args:
+            split (str): ``"train"`` or ``"valid"`` depending on the data split needed.
+            dataset_name (str, optional): name of the dataset to be processed. Defaults to
+                ``"CarperAI/openai_summarize_comparisons"``.
+            max_length (int, optional): maximum length of the dataset sequences.
+                Defaults to 550.
+            root_dir (path, optional): the path where the datasets are stored.
+                Defaults to ``"$HOME/.cache/torchrl/data"``
+            from_disk (bool, optional): if ``True``, :func:`datasets.load_from_disk`
+                will be used. Otherwise, :func:`datasets.load_dataset` will be used.
+                Defaults to ``False``.
+            num_workers (int, optional): number of workers for :meth:`datasets.dataset.map`
+                which is called during tokenization.
+                Defaults to ``max(os.cpu_count() // 2, 1)``.
+
+        Returns: a :class:`PromptData` instance containing a memory-mapped
+            version of the required dataset.
+
+        Examples:
+            >>> data = PromptData.from_dataset("train")
+            >>> print(data)
+            PromptDataTLDR(
+                attention_mask=MemoryMappedTensor(shape=torch.Size([116722, 550]), device=cpu, dtype=torch.int64, is_shared=False),
+                input_ids=MemoryMappedTensor(shape=torch.Size([116722, 550]), device=cpu, dtype=torch.int64, is_shared=False),
+                prompt_rindex=MemoryMappedTensor(shape=torch.Size([116722]), device=cpu, dtype=torch.int64, is_shared=False),
+                labels=MemoryMappedTensor(shape=torch.Size([116722, 550]), device=cpu, dtype=torch.int64, is_shared=False),
+                logits=None,
+                loss=None,
+                batch_size=torch.Size([116722]),
+                device=None,
+                is_shared=False)
+            >>> # data can be sampled from using regular indexing
+            >>> sub_data = data[:3]
+
+        """
+        dataset_name = dataset_name if dataset_name is not None else DEFAULT_DATASET
+        loader = TokenizedDatasetLoader(
+            split,
+            max_length,
+            dataset_name,
+            PromptTensorDictTokenizer,
+            root_dir=root_dir,
+            from_disk=from_disk,
+            num_workers=num_workers,
+        )
+        data = loader.load()
+        return cls(**data, labels=data["input_ids"], batch_size=data.shape)
+
+
+class PromptTensorDictTokenizer(TensorDictTokenizer):
+    """Tokenization recipe for prompt datasets.
+
+    Returns a tokenizer function, which reads an example containing a prompt
+    and a label and tokenizes them.
+
+    Args:
+        tokenizer (tokenizer from transformers library): the tokenizer to use.
+        max_length (int): maximum length of the sequence.
+        key (str, optional): the key where to find the text. Defaults to ``"prompt"``.
+        padding (str, optional): type of padding. Defaults to ``"max_length"``.
+        truncation (bool, optional): whether the sequences should be truncated to max_length.
+        return_tensordict (bool, optional): if ``True``, a TensoDict is returned.
+            Otherwise, a the original data will be returned.
+        device (torch.device, optional): the device where to store the data.
+            This option is ignored if ``return_tensordict=False``.
+
+    The :meth:`__call__` method of this class will execute the following operations:
+
+        - Read the ``prompt`` string contacted with the ``label`` string and tokenize
+          them. The results will be stored in the ``"input_ids"`` TensorDict entry.
+        - Write a ``"prompt_rindex"`` entry with the index of the last valid
+          token from the prompt.
+        - Write a ``"valid_sample"`` which identifies which entry in the
+          tensordict has eough toknens to meet the ``max_length`` criterion.
+        - Return a :class:`tensordict.TensorDict` instance with tokenized inputs.
+
+    The tensordict batch-size will match the batch-size of the input.
+
+    Examples:
+        >>> from transformers import AutoTokenizer
+        >>> tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        >>> tokenizer.pad_token = tokenizer.eos_token
+        >>> example = {
+        ...     "prompt": ["This prompt is long enough to be tokenized.", "this one too!"],
+        ...     "label": ["Indeed it is.", 'It might as well be.'],
+        ... }
+        >>> fn = PromptTensorDictTokenizer(tokenizer, 50)
+        >>> print(fn(example))
+        TensorDict(
+            fields={
+                attention_mask: Tensor(shape=torch.Size([2, 50]), device=cpu, dtype=torch.int64, is_shared=False),
+                input_ids: Tensor(shape=torch.Size([2, 50]), device=cpu, dtype=torch.int64, is_shared=False),
+                prompt_rindex: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False),
+                valid_sample: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.bool, is_shared=False)},
+            batch_size=torch.Size([2]),
+            device=None,
+            is_shared=False)
+
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        max_length,
+        key="prompt",
+        padding="max_length",
+        truncation=True,
+        return_tensordict=True,
+        device=None,
+    ):
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.key = key
+        self.padding = padding
+        self.truncation = truncation
+        self.return_tensordict = return_tensordict
+        self.device = device
+
+    def __call__(self, sample):
+        tokenizer = self.tokenizer
+        max_length = self.max_length
+
+        tokenized_prompts = tokenizer(
+            sample[self.key], max_length=max_length, truncation=True
+        )
+        prompt_rindex = [len(prompt) - 1 for prompt in tokenized_prompts["input_ids"]]
+        tokenized_example = tokenizer(
+            [
+                prompt + label
+                for prompt, label in zip(sample[self.key], sample["label"])
+            ],
+            max_length=max_length,
+            padding=self.padding,
+            truncation=self.truncation,
+        )
+        tokenized_example["prompt_rindex"] = prompt_rindex
+        # drop any examples whose total length when tokenized exceeds block size
+        # with recommended block size of 550, this is only ~0.1% of available examples.
+        # NOTE: to mark as discarded we just save the mask as we cannot change the shape here
+        tokenized_example["valid_sample"] = [True] * len(tokenized_example["input_ids"])
+        for i, input_ids in enumerate(tokenized_example["input_ids"]):
+            if input_ids[-1] != tokenizer.eos_token_id:
+                tokenized_example["valid_sample"][i] = False
+        if self.return_tensordict:
+            return TensorDict.from_dict(dict(tokenized_example), device=self.device)
+        return tokenized_example

torchrl/data/llm/reward.py

@@ -0,0 +1,225 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import importlib
+
+import torch
+from tensordict import tensorclass
+from torchrl.data.llm.dataset import TensorDictTokenizer, TokenizedDatasetLoader
+
+DEFAULT_DATASET = "CarperAI/openai_summarize_comparisons"
+_has_datasets = importlib.util.find_spec("datasets") is not None
+_has_tqdm = importlib.util.find_spec("tqdm") is not None
+
+
+@tensorclass
+class RewardData:
+    """A dataclass for reward model training."""
+
+    input_ids: torch.Tensor
+    attention_mask: torch.Tensor
+    rewards: torch.Tensor | None = None
+    end_scores: torch.Tensor | None = None
+
+
+@tensorclass
+class PairwiseDataset:
+    """Represents a dataset in a pairwise manner (chosen vs rejected).
+
+    Attributes:
+        chosen_data: data to be chosen.
+        rejected_data: corresponding data to be rejected.
+
+    Examples:
+        >>> data = PairwiseDataset.from_dataset("train", max_length=550)
+        >>> print(data)
+        PairwiseDataset(
+            chosen_data=RewardData(
+                attention_mask=MemoryMappedTensor(shape=torch.Size([92534, 550]), device=cpu, dtype=torch.int64, is_shared=False),
+                input_ids=MemoryMappedTensor(shape=torch.Size([92534, 550]), device=cpu, dtype=torch.int64, is_shared=False),
+                rewards=None,
+                end_scores=None,
+                batch_size=torch.Size([92534]),
+                device=None,
+                is_shared=False),
+            rejected_data=RewardData(
+                attention_mask=MemoryMappedTensor(shape=torch.Size([92534, 550]), device=cpu, dtype=torch.int64, is_shared=False),
+                input_ids=MemoryMappedTensor(shape=torch.Size([92534, 550]), device=cpu, dtype=torch.int64, is_shared=False),
+                rewards=None,
+                end_scores=None,
+                batch_size=torch.Size([92534]),
+                device=None,
+                is_shared=False),
+            batch_size=torch.Size([92534]),
+            device=None,
+            is_shared=False)
+
+    """
+
+    chosen_data: RewardData
+    rejected_data: RewardData
+
+    @classmethod
+    def from_dataset(
+        cls,
+        split,
+        dataset_name: str | None = None,
+        max_length: int = 550,
+        root_dir: str | None = None,
+        from_disk: bool = False,
+        num_workers: int | None = None,
+    ):
+        """Returns a :class:`PairwiseDataset` from a dataset name.
+
+        Args:
+            split (str): ``"train"`` or ``"valid"`` depending on the data split needed.
+            dataset_name (str, optional): name of the dataset to be processed. Defaults to
+                ``"CarperAI/openai_summarize_comparisons"``.
+            max_length (int, optional): maximum length of the dataset sequences.
+                Defaults to 550.
+            root_dir (path, optional): the path where the datasets are stored.
+                Defaults to ``"$HOME/.cache/torchrl/data"``
+            from_disk (bool, optional): if ``True``, :func:`datasets.load_from_disk`
+                will be used. Otherwise, :func:`datasets.load_dataset` will be used.
+                Defaults to ``False``.
+
+        Returns: a :class:`PairwiseDataset` instance containing a memory-mapped
+            version of the required dataset.
+
+        Examples:
+            >>> data = PairwiseDataset.from_dataset("train")
+            >>> print(data)
+            PairwiseDataset(
+                chosen_data=RewardData(
+                    attention_mask=MemoryMappedTensor(shape=torch.Size([92534, 550]), device=cpu, dtype=torch.int64, is_shared=False),
+                    input_ids=MemoryMappedTensor(shape=torch.Size([92534, 550]), device=cpu, dtype=torch.int64, is_shared=False),
+                    rewards=None,
+                    end_scores=None,
+                    batch_size=torch.Size([92534]),
+                    device=None,
+                    is_shared=False),
+                rejected_data=RewardData(
+                    attention_mask=MemoryMappedTensor(shape=torch.Size([92534, 550]), device=cpu, dtype=torch.int64, is_shared=False),
+                    input_ids=MemoryMappedTensor(shape=torch.Size([92534, 550]), device=cpu, dtype=torch.int64, is_shared=False),
+                    rewards=None,
+                    end_scores=None,
+                    batch_size=torch.Size([92534]),
+                    device=None,
+                    is_shared=False),
+                batch_size=torch.Size([92534]),
+                device=None,
+                is_shared=False)
+            >>> # data can be sampled from using regular indexing
+            >>> sub_data = data[:3]
+
+        """
+        if dataset_name is None:
+            dataset_name = DEFAULT_DATASET
+        loader = TokenizedDatasetLoader(
+            split,
+            max_length,
+            dataset_name,
+            TensorDictTokenizer,
+            pre_tokenization_hook,
+            root_dir=root_dir,
+            from_disk=from_disk,
+            num_workers=num_workers,
+        )
+        data = loader.load()
+        maxidx = data.shape[0] // 2
+        batch_size = [maxidx]
+        # this is a zero-copy creation, as we index memmap-arrays without
+        # creating new storage.
+        chosen_data = data[:maxidx]
+        rejected_data = data[maxidx:]
+        return cls(
+            chosen_data=RewardData(
+                **chosen_data,
+                batch_size=batch_size,
+            ),
+            rejected_data=RewardData(
+                **rejected_data,
+                batch_size=batch_size,
+            ),
+            batch_size=batch_size,
+        )
+
+
+def pre_tokenization_hook(dataset, min_length=5):
+    """Pre-tokenizer for the reward model (comparison) dataset.
+
+    This function selects all samples where the length of the prompt is
+    sufficient and where the ``"chosen"`` and ``"rejected"`` entries differ.
+
+    Args:
+        dataset (datasets.Dataset): the dataset to process. Should have entries
+            ``"prompt"``, ``"chosen"`` and ``"rejected"``.
+        min_length (int, optional): minimum length of a prompt (in word count).
+
+    Returns: a new ``datasets.Dataset`` with selected prompts under ``"text"``.
+        The first half are the chosen strings and the second the rejected ones,
+        always preceded by the original prompt.
+
+    Examples:
+        >>> from datasets import Dataset
+        >>> data = Dataset.from_dict({
+        ...     "prompt": ["I'm the king"],
+        ...     "chosen": ["It is true, you are the king"],
+        ...     "rejected": ["No, I am the king, you are not"]})
+        >>> print(pre_tokenization_hook(data))
+        Dataset({
+            features: ['text'],
+            num_rows: 2
+        })
+        >>> data = Dataset.from_dict({
+        ...     "prompt": ["I'm the king"],
+        ...     "chosen": ["It is true, you are the king"],
+        ...     "rejected": ["It is true, you are the king"]}) # chosen and rejected match
+        >>> print(pre_tokenization_hook(data))
+        Dataset({
+            features: [],
+            num_rows: 0
+        })
+        >>> data = Dataset.from_dict({
+        ...     "prompt": ["I'm the king"],
+        ...     "chosen": ["Yes"],
+        ...     "rejected": ["No"]}) # chosen and rejected are too short
+        >>> print(pre_tokenization_hook(data))
+        Dataset({
+            features: [],
+            num_rows: 0
+        })
+
+    """
+    if not _has_datasets:
+        raise ImportError(
+            "datasets module couldn't be found. Make sure it is installed in your environment."
+        )
+    from datasets import Dataset as HFDataset
+
+    chosen = []
+    rejected = []
+    if _has_tqdm:
+        from tqdm import tqdm
+
+        pbar = tqdm(dataset)
+    else:
+        pbar = dataset
+    for sample in pbar:
+        prompt = sample["prompt"]
+        chosen_summary = sample["chosen"]
+        rejected_summary = sample["rejected"]
+        if chosen_summary == rejected_summary:
+            continue
+        if (
+            len(chosen_summary.split()) < min_length
+            or len(rejected_summary.split()) < min_length
+        ):
+            continue
+        chosen.append({"text": prompt + "\n" + chosen_summary})
+        rejected.append({"text": prompt + "\n" + rejected_summary})
+
+    return HFDataset.from_list(chosen + rejected)

torchrl/data/llm/topk.py

@@ -0,0 +1,186 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+from collections import defaultdict, deque
+from typing import Any
+
+import torch
+from tensordict import NestedKey, TensorDictBase
+from torchrl._utils import logger as torchrl_logger
+from torchrl.envs.transforms import Transform
+
+
+class TopKRewardSelector(Transform):
+    """A replay-buffer transform that selects the top-k rewards for each prompt.
+
+    Args:
+        total_dialog_turns (int): Number of dialog turns to keep in memory for the top-k selection.
+        topk_size (int): Number of top-k rewards to select. Must be smaller than or equal to total_dialog_turns.
+        prompt_key (NestedKey): Key to the prompt in the tensordict. Defaults to "text".
+        rewards_key (NestedKey): Key to the rewards in the tensordict. Defaults to ("next", "reward").
+        done_key (NestedKey): Key to the done state in the tensordict. Defaults to ("next", "done").
+        verbose (bool): Whether to print verbose information. Defaults to `False`.
+
+    Example:
+        >>> from torchrl.data import ReplayBuffer, LazyStackStorage, SamplerWithoutReplacement
+        >>> from tensordict import TensorDict, lazy_stack
+        >>> import torch
+        >>> from torchrl.data.llm.topk import TopKRewardSelector
+        >>> # Create a replay buffer with 50 items, a sampler that samples without replacement, and a batch size of 5
+        >>> rb = ReplayBuffer(
+        ...     storage=LazyStackStorage(50),
+        ...     sampler=SamplerWithoutReplacement,
+        ...     batch_size=5,
+        ... )
+        >>> # Create a tensordict with 50 items, each with 10 dialog turns
+        >>> td = lazy_stack(
+        ...     [
+        ...         TensorDict(
+        ...             {
+        ...                 ("next", "done"): torch.full((1, 1), True),
+        ...                 # Reward for i+5 tokens
+        ...                 ("next", "reward"): torch.full((i + 5, 1), i),
+        ...                 # total of 10 dialogs per prompt
+        ...                 "text": f"Prompt {i // 5}",
+        ...             }
+        ...         )
+        ...         for i in range(50)
+        ...     ]
+        ... )
+        >>> # Create a top-k reward selector with 5 dialog turns and a top-k size of 3
+        >>> topk = TopKRewardSelector(total_dialog_turns=5, topk_size=3)
+        >>> rb.append_transform(topk)
+        >>> for _td in td.chunk(25):
+        ...     rb.extend(_td)
+        >>> # Only wrote top3 of 50 items in 10 groups of 5
+        >>>  assert rb.write_count == 30
+        >>> assert len(rb) == 30
+        >>> r3 = rb[:3].get(("next", "reward"), as_padded_tensor=True).squeeze()
+        >>> # 0 and 1 are missing because they're not part of the top-k
+        >>> assert (
+        ...     r3 == torch.tensor(
+        ...         [
+        ...             [4, 4, 4, 4, 4, 4, 4, 4, 4],
+        ...             [3, 3, 3, 3, 3, 3, 3, 3, 0],
+        ...             [2, 2, 2, 2, 2, 2, 2, 0, 0],
+        ...         ]
+        ...     )
+        ... ).all()
+    """
+
+    def __init__(
+        self,
+        total_dialog_turns: int,
+        topk_size: int,
+        prompt_key: NestedKey = ("text", "prompt"),
+        rewards_key: NestedKey = ("next", "reward"),
+        done_key: NestedKey = ("next", "done"),
+        verbose: bool = True,
+    ):
+        super().__init__()
+        self.in_keys = [prompt_key, rewards_key, done_key]
+        self.prompt_key = prompt_key
+        self.rewards_key = rewards_key
+        self.done_key = done_key
+        self.queues = defaultdict(lambda: deque(maxlen=total_dialog_turns))
+        self.total_dialog_turns = total_dialog_turns
+        self.topk_size = topk_size
+        if topk_size > total_dialog_turns:
+            raise ValueError(
+                f"topk_size must be smaller than or equal to total_dialog_turns, got {topk_size=} and {total_dialog_turns=}"
+            )
+        self.verbose = verbose
+
+    def forward(self, tensordict: TensorDictBase) -> Any:
+        return tensordict
+
+    def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase:
+        # Tensordict can be any number of dims, but it must contain entire trajectories
+        if tensordict.ndim == 1:
+            # Check how many done states we have
+            num_done = tensordict[self.done_key].sum()
+            if num_done > 1:
+                done_idx = tensordict[self.done_key].nonzero(as_tuple=True)[0] + 1
+                splits = torch.cat([done_idx.new_zeros((1,)), done_idx], dim=0).diff()
+                tensordicts = tensordict.split(splits)
+                tensordicts = [self._inv_call(td) for td in tensordicts]
+                tensordicts = [td for td in tensordicts if td is not None]
+                return torch.cat(tensordicts) if tensordicts else None
+            # Then we have a single trajectory
+            if not tensordict[-1][self.done_key].all():
+                raise RuntimeError("Expected the trajectory to be done.")
+            prompt = tensordict[0][self.prompt_key]
+            if not isinstance(prompt, str):
+                raise TypeError(f"Expected a string as prompt, got {type(prompt)=}")
+            self.queues[prompt].append(tensordict)
+            if len(self.queues[prompt]) == self.total_dialog_turns:
+                if self.verbose:
+                    torchrl_logger.info(f"Getting top-k rewards for {prompt=}")
+                # Cat is the most robust way to combine the trajs
+                tds = torch.cat(list(self.queues[prompt]), -1)
+                # Collect rewards
+                reward = tds.get(self.rewards_key, as_nested_tensor=True)
+                reward = self._aggregate_rewards(reward)
+                # Check if all rewards are equal
+                if (reward == reward[0]).all():
+                    # If all rewards are equal, we can't select top-k
+                    if self.verbose:
+                        torchrl_logger.warning(
+                            f"All rewards are equal ({reward.unique()=})"
+                        )
+                    return
+                # Filter out rewards below median
+                median_reward = reward.median(dim=-1, keepdim=True)[0]
+                mask = reward > median_reward
+                filtered_reward = reward[mask]
+                filtered_indices = mask.nonzero(as_tuple=True)[0]
+                # Get top-k from filtered rewards
+                topk_reward = filtered_reward.topk(
+                    k=min(self.topk_size, len(filtered_indices)), dim=-1
+                )
+                if not topk_reward.indices.numel():
+                    if self.verbose:
+                        torchrl_logger.warning(
+                            f"No top-{self.topk_size} rewards found ({reward=})"
+                        )
+                    return
+                # Map back to original indices
+                selected_indices = filtered_indices[topk_reward.indices]
+                tds = tds[selected_indices]
+                if self.verbose:
+                    torchrl_logger.info(
+                        f"Selected top-{self.topk_size} rewards, with reward {topk_reward.values=}"
+                    )
+                return tds
+            return
+        elif tensordict.ndim > 2:
+            # keep the time dim at the end
+            tensordict = tensordict.flatten(0, -2)
+        trajs = tensordict.unbind(-1)
+        # Iterate over the trajectories
+        result = []
+        for traj in trajs:
+            td_out = self._inv_call(traj)
+            if td_out is None:
+                continue
+            result.append(td_out)
+        if result:
+            return torch.cat(result, -1)
+        return
+
+    def _aggregate_rewards(self, reward: torch.Tensor) -> torch.Tensor:
+        """Aggregate the rewards across the dialog turns.
+
+        `reward` is expected to be a nested tensor.
+
+        The default implementation is to take the mean of the rewards across the dialog turns.
+        """
+        # reward = reward.to_padded_tensor(padding=0.0)
+        if reward.ndim < 2 or reward.ndim > 3:
+            raise ValueError(
+                f"Expected reward to be a 2D or 3D tensor, got {reward.ndim}D tensor"
+            )
+        return reward.mean(dim=-2).squeeze(-1)