Update

[ghstack-poisoned]

Author: Vincent Moens

torchrl/data/postprocs/postprocs.py

@@ -12,7 +12,6 @@
 from torch import nn
 
 
-
 def _get_reward(
     gamma: float,
     reward: torch.Tensor,
@@ -367,6 +366,7 @@ def __init__(
         discount: float = 1.0,
     ):
         from torchrl.objectives.value.functional import reward2go
+
         super().__init__()
         self.in_keys = [unravel_key(reward_key), unravel_key(done_key)]
         if reward_key_out is None:

torchrl/envs/custom/llm.py

@@ -80,7 +80,9 @@ def __init__(
             self._batch_locked = False
         else:
             self._batch_locked = True
-        super().__init__(device=device, batch_size=() if batch_size is None else (batch_size,))
+        super().__init__(
+            device=device, batch_size=() if batch_size is None else (batch_size,)
+        )
         self.str2str = str2str
         self.vocab_size = vocab_size
         self.observation_key = unravel_key(token_key)
@@ -92,29 +94,21 @@ def __init__(
         # self.action_key = unravel_key(action_key)
         if str2str:
             self.full_observation_spec_unbatched = Composite(
-                {
-                    token_key: NonTensor(
-                        example_data="a string", batched=True, shape=()
-                    )
-                }
+                {token_key: NonTensor(example_data="a string", batched=True, shape=())}
             )
             self.full_action_spec_unbatched = Composite(
                 {action_key: NonTensor(example_data="a string", batched=True, shape=())}
             )
         else:
             if vocab_size is None:
                 observation_spec = {
-                        token_key: Unbounded(
-                            shape=(-1,), dtype=torch.int64, device=device
-                        )
-                    }
+                    token_key: Unbounded(shape=(-1,), dtype=torch.int64, device=device)
+                }
                 if attention_key is not None:
                     observation_spec[attention_key] = Unbounded(
-                            shape=(-1,), dtype=torch.int64, device=device
-                        )
-                self.full_observation_spec_unbatched = Composite(
-                    observation_spec
-                )
+                        shape=(-1,), dtype=torch.int64, device=device
+                    )
+                self.full_observation_spec_unbatched = Composite(observation_spec)
                 self.full_action_spec_unbatched = Composite(
                     {
                         action_key: Unbounded(
@@ -325,7 +319,13 @@ def _make_next_obs(
             if self.attention_key is not None:
                 attention_mask = tensordict.get(self.attention_key)
                 n = action.shape[-1] - attention_mask.shape[-1]
-                attention_mask = torch.cat([attention_mask, attention_mask.new_ones(attention_mask.shape[:-1] + (n,))], -1)
+                attention_mask = torch.cat(
+                    [
+                        attention_mask,
+                        attention_mask.new_ones(attention_mask.shape[:-1] + (n,)),
+                    ],
+                    -1,
+                )
                 nex_td.set(self.attention_key, attention_mask)
             return nex_td
 
@@ -384,7 +384,7 @@ def _make_next_obs(
 
     def _reset(self, tensordict: TensorDictBase, **kwargs) -> TensorDictBase:
         # We should have an observation by this time, if not raise an exception
-        print('tensordict', tensordict)
+        print("tensordict", tensordict)
         if tensordict is None or self.observation_key not in tensordict.keys(
             isinstance(self.observation_key, tuple)
         ):

torchrl/envs/transforms/rlhf.py

@@ -461,7 +461,7 @@ def _load_from_dataloader(self, reset: torch.Tensor | None = None):
             raise ValueError(
                 f"Unrecognized data type: {type(data)} with keys {self.data_keys}."
             )
-        print('out', out)
+        print("out", out)
         if self.use_buffer:
             if not out.ndim:
                 out = out.unsqueeze(0)

torchrl/modules/llm/transformers.py

@@ -5,11 +5,17 @@
 
 # TODO: lazy imports
 
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from tensordict.nn import TensorDictModule as Mod, TensorDictSequential as Seq, TensorDictModuleBase, WrapModule
-from tensordict import NestedKey, TensorDictBase, TensorDict
-import transformers
 import torch
+import transformers
+from tensordict import NestedKey, TensorDict, TensorDictBase
+from tensordict.nn import (
+    TensorDictModule as Mod,
+    TensorDictModuleBase,
+    TensorDictSequential as Seq,
+    WrapModule,
+)
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
 
 def _maybe_clear_device(td):
     if td.device is None:
@@ -30,7 +36,9 @@ def log_probs_from_scores(td: TensorDictBase) -> TensorDictBase:
     # TODO: how do we avoid getting these?
     del td["tokens_out", "past_key_values"]
     scores = dict(td["tokens_out", "scores"].items())
-    scores = torch.stack([scores[str(k)] for k in range(len(scores))], 1)  # shape (B, seq-len, vocab_size)
+    scores = torch.stack(
+        [scores[str(k)] for k in range(len(scores))], 1
+    )  # shape (B, seq-len, vocab_size)
     logits = scores - scores.logsumexp(dim=-1, keepdim=True)
     td["logits"] = scores
     del td["tokens_out", "scores"]
@@ -40,33 +48,34 @@ def log_probs_from_scores(td: TensorDictBase) -> TensorDictBase:
     td["log_probs"] = log_probs
     return td
 
+
 def log_probs_from_logits(td: TensorDictBase) -> TensorDictBase:
     # TODO: how do we avoid getting these?
     del td["forward", "past_key_values"]
     scores = td["forward", "logits"]
     logits = scores - scores.logsumexp(dim=-1, keepdim=True)
     td["logits"] = scores
     del td["forward"]
-    seq_len = scores.shape[1]
+    scores.shape[1]
     tokens = td["tokens_in", "input_ids"]
     log_probs = logits.gather(-1, tokens.unsqueeze(-1))
     td["log_probs"] = log_probs
     return td
 
 
 def from_hf_transformers(
-        model: transformers.modeling_utils.PreTrainedModel,
-        *,
-        generate: bool = True,
-        return_log_probs: bool = True,
-        tokenizer: transformers.tokenization_utils.PreTrainedTokenizer | None = None,
-        from_text: bool = False,
-        device: torch.device | None = None,
-        text_key: NestedKey = "text",
-        input_key: NestedKey = "input_ids",
-        kwargs: dict | None = None,
-        tokenizer_kwargs: dict | None = None,
-        ) -> TensorDictModuleBase:
+    model: transformers.modeling_utils.PreTrainedModel,
+    *,
+    generate: bool = True,
+    return_log_probs: bool = True,
+    tokenizer: transformers.tokenization_utils.PreTrainedTokenizer | None = None,
+    from_text: bool = False,
+    device: torch.device | None = None,
+    text_key: NestedKey = "text",
+    input_key: NestedKey = "input_ids",
+    kwargs: dict | None = None,
+    tokenizer_kwargs: dict | None = None,
+) -> TensorDictModuleBase:
 
     # TODO: Seq should have a return_log_prob and be of ProbabilisticTDSequential type for instance checks
 
@@ -98,7 +107,7 @@ def from_hf_transformers(
             lambda tensor: tensor.to(device),
             in_keys=["tokens_in"],
             out_keys=["tokens_in"],
-            strict=True
+            strict=True,
         )
 
     if generate:
@@ -109,7 +118,10 @@ def from_hf_transformers(
                 raise RuntimeError
         if not kwargs.setdefault("return_dict_in_generate", True):
             raise RuntimeError
-        if kwargs.setdefault("tokenizer", tokenizer) is not tokenizer and tokenizer is not None:
+        if (
+            kwargs.setdefault("tokenizer", tokenizer) is not tokenizer
+            and tokenizer is not None
+        ):
             raise RuntimeError
 
         module_dict["generate"] = Mod(
@@ -128,8 +140,8 @@ def from_hf_transformers(
             module_dict["extract_log_probs"] = WrapModule(
                 log_probs_from_scores,
                 in_keys=[("tokens_out", "sequences"), ("tokens_out", "scores")],
-                out_keys=["logits", "log_probs"]
-                )
+                out_keys=["logits", "log_probs"],
+            )
         if from_text:
             module_dict["decode"] = Mod(
                 tokenizer.batch_decode,
@@ -159,8 +171,8 @@ def from_hf_transformers(
         module_dict["extract_log_probs"] = WrapModule(
             log_probs_from_logits,
             in_keys=[("tokens_in", "input_ids"), ("forward", "logits")],
-            out_keys=["logits", "log_probs"]
-            )
+            out_keys=["logits", "log_probs"],
+        )
     if device:
         module_dict["to_source_device"] = _maybe_set_device
     return Seq(module_dict)
@@ -171,16 +183,18 @@ def from_hf_transformers(
     model_name = "Qwen/Qwen2.5-7B-Instruct"
 
     model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype="auto",
-        device_map="auto"
+        model_name, torch_dtype="auto", device_map="auto"
     )
     tokenizer = AutoTokenizer.from_pretrained(model_name)
 
     tokenizer.padding_side = "left"
 
-    m = from_hf_transformers(model, tokenizer=tokenizer, from_text=True, device="cuda:0", generate=True)
+    m = from_hf_transformers(
+        model, tokenizer=tokenizer, from_text=True, device="cuda:0", generate=True
+    )
     td = m(TensorDict(text="a text"))
 
-    m = from_hf_transformers(model, tokenizer=tokenizer, from_text=True, device="cuda:0", generate=False)
+    m = from_hf_transformers(
+        model, tokenizer=tokenizer, from_text=True, device="cuda:0", generate=False
+    )
     td = m(TensorDict(text="a text"))