Update

[ghstack-poisoned]

Author: Vincent Moens

docs/source/reference/data.rst

@@ -1133,6 +1133,9 @@ efficient sampling.
     get_dataloader
     ConstantKLController
     AdaptiveKLController
+    LLMData
+    LLMInput
+    LLMOutput
 
 
 Utils

examples/rlhf/data/__init__.py

@@ -1,3 +1,3 @@
-from torchrl.data.rlhf.prompt import get_prompt_dataloader_tldr
+from torchrl.data.llm.prompt import get_prompt_dataloader_tldr
 
 __all__ = ["get_prompt_dataloader_tldr"]

examples/rlhf/models/reward.py

@@ -8,7 +8,7 @@
 from tensordict.nn import TensorDictModule
 from torchrl._utils import logger as torchrl_logger
 
-from torchrl.modules.models.rlhf import GPT2RewardModel
+from torchrl.modules.models.llm import GPT2RewardModel
 
 
 def init_reward_model(

examples/rlhf/train.py

@@ -17,8 +17,8 @@
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from torchrl._utils import logger as torchrl_logger
 
-from torchrl.data.rlhf.dataset import get_dataloader
-from torchrl.data.rlhf.prompt import PromptData
+from torchrl.data.llm.dataset import get_dataloader
+from torchrl.data.llm.prompt import PromptData
 from utils import get_file_logger, resolve_name_or_path, setup
 
 

examples/rlhf/train_reward.py

@@ -9,8 +9,8 @@
 from models.reward import init_reward_model
 from torch.optim.lr_scheduler import CosineAnnealingLR
 from torchrl._utils import logger as torchrl_logger
-from torchrl.data.rlhf.dataset import get_dataloader
-from torchrl.data.rlhf.reward import PairwiseDataset
+from torchrl.data.llm.dataset import get_dataloader
+from torchrl.data.llm.reward import PairwiseDataset
 from utils import get_file_logger, resolve_name_or_path, setup
 
 

examples/rlhf/train_rlhf.py

@@ -6,7 +6,7 @@
 import hydra
 import torch
 from models.actor_critic import init_actor_critic
-from torchrl.data.rlhf.utils import AdaptiveKLController, RolloutFromModel
+from torchrl.data.llm.utils import AdaptiveKLController, RolloutFromModel
 
 from torchrl.record.loggers import get_logger
 

examples/rlhf/utils.py

@@ -22,9 +22,9 @@
     TensorDictReplayBuffer,
     TensorStorage,
 )
+from torchrl.data.llm.dataset import get_dataloader
+from torchrl.data.llm.prompt import PromptData
 from torchrl.data.replay_buffers import SamplerWithoutReplacement
-from torchrl.data.rlhf.dataset import get_dataloader
-from torchrl.data.rlhf.prompt import PromptData
 from torchrl.objectives import ClipPPOLoss
 from torchrl.objectives.value import GAE
 

test/assets/generate.py

@@ -5,6 +5,7 @@
 
 """Script used to generate the mini datasets."""
 import multiprocessing as mp
+import pathlib
 
 try:
     mp.set_start_method("spawn")
@@ -14,8 +15,8 @@
 
 from datasets import Dataset, DatasetDict, load_dataset
 
-from torchrl.data.rlhf.dataset import get_dataloader
-from torchrl.data.rlhf.prompt import PromptData
+from torchrl.data.llm.dataset import get_dataloader
+from torchrl.data.llm.prompt import PromptData
 
 
 def generate_small_dataset(comparison=True):
@@ -42,7 +43,7 @@ def get_minibatch():
             batch_size=16,
             block_size=33,
             tensorclass_type=PromptData,
-            dataset_name="../datasets_mini/openai_summarize_tldr",
+            dataset_name=f"{pathlib.Path(__file__).parent}/../datasets_mini/openai_summarize_tldr",
             device="cpu",
             num_workers=2,
             infinite=False,
@@ -52,9 +53,12 @@ def get_minibatch():
             root_dir=tmpdir,
         )
         for data in dl:
-            data = data.clone().memmap_("test/datasets_mini/tldr_batch/")
+            data = data.clone().memmap_(
+                f"{pathlib.Path(__file__).parent}/../datasets_mini/tldr_batch/"
+            )
             break
 
 
 if __name__ == "__main__":
+    generate_small_dataset(False)
     get_minibatch()

test/assets/tldr_batch.zip

@@ -14,7 +14,7 @@
 
 from torch import distributions as dist, nn
 from torchrl.data import Binary, Bounded, Categorical, Composite, MultiOneHot, OneHot
-from torchrl.data.rlhf.dataset import _has_transformers
+from torchrl.data.llm.dataset import _has_transformers
 from torchrl.modules import MLP, SafeModule, TanhDelta, TanhNormal
 from torchrl.modules.tensordict_module.actors import (
     _process_action_space_spec,