Update (base update)

[ghstack-poisoned]

Author: vmoens

torchrl/envs/transforms/llm.py

@@ -10,13 +10,7 @@
 from typing import Any, Callable, Iterable, Literal
 
 import torch
-from tensordict import (
-    lazy_stack,
-    NestedKey,
-    TensorDict,
-    TensorDictBase,
-    unravel_key,
-)
+from tensordict import lazy_stack, NestedKey, TensorDict, TensorDictBase, unravel_key
 from tensordict.nn import ProbabilisticTensorDictModule, TensorDictParams
 from tensordict.utils import _zip_strict, is_seq_of_nested_key
 from torch import nn

torchrl/modules/llm/vllm_policy.py

@@ -244,11 +244,10 @@ def to_list(tokens, attention_mask):
     )
 
     if generate_kwargs is None:
-        generate_kwargs = {
-            "detokenize": False,
-            "prompt_logprobs": not generate,
-            "logprobs": return_log_probs,
-        }
+        generate_kwargs = {}
+    generate_kwargs.setdefault("detokenize", False)
+    generate_kwargs.setdefault("prompt_logprobs", not generate)
+    generate_kwargs.setdefault("logprobs", return_log_probs)
     if not generate:
         generate_kwargs["max_tokens"] = 1
     sampling_params = SamplingParams(**generate_kwargs)