torchrl 0.11.0__cp314-cp314-win_amd64.whl

torchrl/data/llm/history.py

@@ -0,0 +1,1378 @@
+# 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 dataclasses
+
+import re
+from typing import Literal, TYPE_CHECKING
+
+import torch
+
+from tensordict import (
+    lazy_stack,
+    LazyStackedTensorDict,
+    list_to_stack,
+    TensorClass,
+    TensorDict,
+)
+from tensordict.utils import _maybe_correct_neg_dim
+from torchrl._utils import logger as torchrl_logger
+
+if TYPE_CHECKING:
+    import transformers
+
+
+# Global storage for custom templates and their metadata
+_CHAT_TEMPLATES = {
+    "chatml_format": """{% for message in messages %}
+    {%- if message['role'] == 'assistant' %}
+    {% generation %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endgeneration %}
+    {%- else %}
+    {{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}
+    {%- endif %}
+{% endfor %}
+{%- if add_generation_prompt %}
+    {% generation %}{{- '<|im_start|>assistant\n' }}{% endgeneration %}
+{%- endif %}
+""",
+    "qwen": """
+{%- if tools %}
+    {{- '<|im_start|>system\\n' }}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- messages[0]['content'] }}
+    {%- else %}
+        {{- 'You are a helpful assistant.' }}
+    {%- endif %}
+    {{- "\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n" }}
+{%- else %}
+    {%- if messages[0]['role'] == 'system' %}
+        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}
+    {%- else %}
+        {{- '<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n' }}
+    {%- endif %}
+{%- endif %}
+{%- for message in messages %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}
+    {%- elif (message.role == "assistant" and not message.tool_calls) %}
+    {% generation %}    {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}    {% endgeneration %}
+    {%- elif message.role == "assistant" %}
+        {% generation %}{{- '<|im_start|>' + message.role }}
+        {%- if message.content %}
+            {{- '\\n' + message.content }}
+        {%- endif %}
+        {%- for tool_call in message.tool_calls %}
+            {%- if tool_call.function is defined %}
+                {%- set tool_call = tool_call.function %}
+            {%- endif %}
+            {{- '\\n<tool_call>\\n{\\\"name\\\": \\\"' }}
+            {{- tool_call.name }}
+            {{- '\\\", \\\"arguments\\\": ' }}
+            {{- tool_call.arguments | tojson }}
+            {{- '}\\n</tool_call>' }}
+        {%- endfor %}
+        {{- '<|im_end|>\\n' }}{% endgeneration %}
+    {%- elif message.role == "tool" %}
+        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>tool' }}
+        {%- endif %}
+        {{- '\\n<tool_response>\\n' }}
+        {%- if message.tool_responses %}
+            {{- message.tool_responses }}
+        {%- else %}
+            {{- message.content }}
+        {%- endif %}
+        {{- '\\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {% generation %}{{- '<|im_start|>assistant\\n' }}{% endgeneration %}
+{%- endif %}
+""",
+    "dialogpt": """{% for message in messages %}{% if message['role'] == 'assistant' %}{% generation %}{{ message['content'] }}{% endgeneration %}{{ eos_token }}{% elif message['role'] == 'user' %}{{ message['content'] }}{{ eos_token }}{% endif %}{% endfor %}{% if add_generation_prompt %}{% generation %}{{ ' ' }}{% endgeneration %}{% endif %}""",
+    "falcon": """{% for message in messages %}{% if message['role'] == 'assistant' %}{% generation %}{{ 'Assistant: ' + message['content'] }}{% endgeneration %}\n\n{% elif message['role'] == 'user' %}{{ 'User: ' + message['content'] }}\n\n{% elif message['role'] == 'system' %}{{ message['content'] }}\n\n{% endif %}{% endfor %}{% if add_generation_prompt %}{% generation %}{{ 'Assistant: ' }}{% endgeneration %}{% endif %}""",
+    "deepseek": """{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{{ bos_token }}{% for message in messages %}{% if message['role'] == 'user' %}{{ 'User: ' + message['content'] + '\n\n' }}{% elif message['role'] == 'assistant' %}{% generation %}{{ 'Assistant: ' + message['content'] + eos_token }}{% endgeneration %}{% elif message['role'] == 'system' %}{{ message['content'] + '\n\n' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{% generation %}{{ 'Assistant:' }}{% endgeneration %}{% endif %}""",
+    "llama": """{{- bos_token }}
+{%- if messages[0]['role'] == 'system' %}
+    {%- set system_message = messages[0]['content']|trim %}
+    {%- set messages = messages[1:] %}
+{%- else %}
+    {%- set system_message = "" %}
+{%- endif %}
+{%- if system_message %}
+    {{- "<|header_start|>system<|header_end|>\n\n" }}
+    {{- system_message }}
+    {{- "<|eot|>" }}
+{%- endif %}
+{%- for message in messages %}
+    {%- if message['role'] == 'assistant' %}
+    {% generation %}{{- '<|header_start|>' + message['role'] + '<|header_end|>\n\n' }}
+        {%- if message['content'] is string %}
+            {{- message['content'] }}
+        {%- else %}
+            {%- for content in message['content'] %}
+                {%- if content['type'] == 'text' %}
+                    {{- content['text'] | trim }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+    {{- "<|eot|>" }}{% endgeneration %}
+    {%- else %}
+    {{- '<|header_start|>' + message['role'] + '<|header_end|>\n\n' }}
+        {%- if message['content'] is string %}
+            {{- message['content'] }}
+        {%- else %}
+            {%- for content in message['content'] %}
+                {%- if content['type'] == 'text' %}
+                    {{- content['text'] | trim }}
+                {%- endif %}
+            {%- endfor %}
+        {%- endif %}
+    {{- "<|eot|>" }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {% generation %}{{- '<|header_start|>assistant<|header_end|>\n\n' }}{% endgeneration %}
+{%- endif %}""",
+}
+
+# Global storage for custom template metadata
+_CUSTOM_INVERSE_PARSERS = {}
+_CUSTOM_MODEL_FAMILY_KEYWORDS = {}
+
+
+def add_chat_template(
+    template_name: str,
+    template: str,
+    inverse_parser: callable | None = None,
+    model_family_keywords: list[str] | None = None,
+) -> None:
+    r"""Add a custom chat template to the global template dictionary.
+
+    This function allows you to add custom chat templates for new model families
+    that support assistant token masking via the `{% generation %}` keyword.
+
+    Args:
+        template_name (str): The name of the template (e.g., "llama", "mistral").
+            This name will be used in the `chat_template_name` parameter of
+            `History.apply_chat_template()` and `History.from_text()`.
+        template (str): The Jinja2 template string. Must include `{% generation %}`
+            blocks around assistant message content to enable token masking.
+        inverse_parser (callable, optional): A function that parses formatted text back
+            into a History object. Should have signature `(text: str) -> History`.
+            If None, a basic parser will be used.
+        model_family_keywords (list[str], optional): Keywords to detect this model family
+            in the auto-detection logic. For example, ["llama", "meta-llama"] for Llama models.
+            If provided, the template will be automatically selected for models containing
+            these keywords in their name.
+
+    Example:
+        >>> from torchrl.data.llm.chat import add_chat_template, History
+        >>> from transformers import AutoTokenizer
+        >>>
+        >>> # Add a custom template for Llama models
+        >>> llama_template = '''
+        ... {% for message in messages %}
+        ... {%- if message['role'] == 'user' %}
+        ... {{ '<s>[INST] ' + message['content'] + ' [/INST]' }}
+        ... {%- elif message['role'] == 'assistant' %}
+        ... {% generation %}{{ message['content'] + '</s>' }}{% endgeneration %}
+        ... {%- endif %}
+        ... {% endfor %}
+        ... {%- if add_generation_prompt %}
+        ... {% generation %}{{ ' ' }}{% endgeneration %}
+        ... {%- endif %}
+        ... '''
+        >>>
+        >>> def parse_llama_text(text: str) -> History:
+        ...     # Custom parser for Llama format
+        ...     import re
+        ...     pattern = r'<s>\[INST\]\s*(.*?)\s*\[/INST\]\s*(.*?)</s>'
+        ...     matches = re.findall(pattern, text, re.DOTALL)
+        ...     messages = []
+        ...     for user_content, assistant_content in matches:
+        ...         messages.append(History(role="user", content=user_content.strip()))
+        ...         messages.append(History(role="assistant", content=assistant_content.strip()))
+        ...     return lazy_stack(messages)
+        >>>
+        >>> # Add the template with auto-detection
+        >>> add_chat_template(
+        ...     template_name="llama",
+        ...     template=llama_template,
+        ...     inverse_parser=parse_llama_text,
+        ...     model_family_keywords=["llama", "meta-llama"]
+        ... )
+        >>>
+        >>> # Now you can use it with auto-detection
+        >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
+        >>> history = History.from_chats([[
+        ...     {"role": "user", "content": "Hello"},
+        ...     {"role": "assistant", "content": "Hi there!"}
+        ... ]])
+        >>>
+        >>> # Auto-detection will use the llama template
+        >>> result = history.apply_chat_template(
+        ...     tokenizer=tokenizer,
+        ...     add_generation_prompt=False,
+        ...     return_dict=True,
+        ...     return_assistant_tokens_mask=True,
+        ... )
+        >>>
+        >>> # Or use it explicitly
+        >>> result = history.apply_chat_template(
+        ...     tokenizer=tokenizer,
+        ...     chat_template_name="llama",
+        ...     add_generation_prompt=False,
+        ...     return_dict=True,
+        ...     return_assistant_tokens_mask=True,
+        ... )
+
+    .. note:
+        - The template must include `{% generation %}` blocks around assistant message
+          content to enable assistant token masking.
+        - The inverse parser should handle the specific format of your template.
+        - Model family keywords are case-insensitive and matched against the tokenizer's
+          `name_or_path` attribute.
+        - Templates are stored globally and persist for the duration of the Python session.
+    """
+    global _CHAT_TEMPLATES, _CUSTOM_INVERSE_PARSERS, _CUSTOM_MODEL_FAMILY_KEYWORDS  # noqa: F824
+
+    # Validate template contains generation blocks
+    if "{% generation %}" not in template:
+        raise ValueError(
+            f"Template '{template_name}' must include '{{% generation %}}' blocks "
+            "around assistant message content to enable token masking."
+        )
+
+    # Add template to dictionary
+    _CHAT_TEMPLATES[template_name] = template
+
+    # Store inverse parser if provided
+    if inverse_parser is not None:
+        _CUSTOM_INVERSE_PARSERS[template_name] = inverse_parser
+
+    # Store model family keywords if provided
+    if model_family_keywords is not None:
+        _CUSTOM_MODEL_FAMILY_KEYWORDS[template_name] = model_family_keywords
+
+    torchrl_logger.info(
+        f"Added custom chat template '{template_name}' with assistant token masking support"
+    )
+
+
+# We need the 'shadow' flag to avoid having tensordict complaining about 'type'/'size' etc. fields
+class ContentBase(TensorClass["nocast", "shadow"]):
+    """Base class for all message content types.
+
+    Attributes:
+        type (str): The type of the content.
+        text (str, optional): The text content.
+        url (str, optional): The URL content.
+        data (str, optional): The data content.
+        mime_type (str, optional): The MIME type of the content.
+        name (str, optional): The name of the content.
+        size (int, optional): The size of the content.
+        function_name (str, optional): The name of the function.
+        function_args (dict, optional): The arguments of the function.
+
+    Examples:
+        >>> from tensordict import lazy_stack
+        >>> content1 = ContentBase(type="text", text="Hello, world!")
+        >>> print(content1)
+        ContentBase(
+            text=NonTensorData(data=Hello, world!, batch_size=torch.Size([]), device=None),
+            type=NonTensorData(data=text, batch_size=torch.Size([]), device=None),
+            url=None,
+            data=None,
+            mime_type=None,
+            name=None,
+            size=None,
+            function_name=None,
+            function_args=None,
+            batch_size=torch.Size([]),
+            device=None,
+            is_shared=False)
+        >>> content2 = ContentBase(type="image", url="https://example.com/image.jpg")
+        >>> print(content2)
+        ContentBase(
+            type=NonTensorData(data=image, batch_size=torch.Size([]), device=None),
+            url=NonTensorData(data=https://example.com/image.jpg, batch_size=torch.Size([]), device=None),
+            text=None,
+            data=None,
+            mime_type=None,
+            name=None,
+            size=None,
+            function_name=None,
+            function_args=None,
+            batch_size=torch.Size([]),
+            device=None,
+            is_shared=False)
+        >>> content = lazy_stack([content1, content2])
+        >>> print(content)
+        ContentBase(
+            type=NonTensorStack(
+                ['text', 'image'],
+                batch_size=torch.Size([2]),
+                device=None),
+            url=None,
+            data=None,
+            mime_type=None,
+            name=None,
+            size=None,
+            function_name=None,
+            function_args=None,
+            text=None,
+            batch_size=torch.Size([2]),
+            device=None,
+            is_shared=False)
+        >>> # A content is typically used in a History object. Usually, its batch dimension is
+        >>> #  one dimension greater than the History object.
+        >>> history = History(role="user", content=content)
+
+    """
+
+    type: Literal[
+        "text", "image", "audio", "video", "file", "function_call"
+    ]  # Required: "text", "image", "audio", "video", "file", "function_call"
+
+    # Text content
+    text: str | None = None
+
+    # Media/file content (either URL or data)
+    url: str | None = None  # HTTP URL to content
+    data: str | None = None  # Base64 encoded content
+
+    # Metadata
+    mime_type: str | None = None  # "image/jpeg", "audio/mp3", "application/pdf"
+    name: str | None = None  # Original filename or description
+    size: int | None = None  # File size in bytes
+
+    # Function calling (for AI agents)
+    function_name: str | None = None
+    function_args: dict | None = None
+
+
+class History(TensorClass["nocast"]):
+    """A class representing a structured history of messages in a conversation, designed for efficient manipulation and integration with language models.
+
+    The `History` class provides a centralized API for managing conversational data, offering several advantages over
+    traditional list-based approaches:
+
+    - Centralized API for conversion to and from string formats, facilitating seamless integration with language models.
+    - Efficient methods to append, extend, and reshape history elements, enabling dynamic construction of conversation
+      trajectories, especially useful in reinforcement learning environments.
+    - Interoperability with the `transformers` API, allowing for easy tokenization and preparation of input data.
+    - **Assistant token masking support** across multiple model families for reinforcement learning applications.
+
+    **Recent Changes:**
+    - **ChatHistory Integration**: History objects are now used within :class:`~torchrl.modules.llm.policies.ChatHistory`
+      containers for structured conversation management in LLM environments.
+    - **Modular Wrapper Support**: Both vLLMWrapper and TransformersWrapper now use History objects when `input_mode="history"`
+      is specified, providing consistent conversation state management.
+    - **Environment Integration**: ChatEnv and related environments use History objects for state management and conversation tracking.
+
+    .. note:: The `"<none>"` role is used to indicate that the element is a placeholder,
+        for example when the tool call was not executed but a stack requires a certain number of elements
+        per batch to have congruent shapes. The :meth:`~torchrl.data.llm.chat.History.apply_chat_template`
+        method will remove the `<none>` role from the history.
+
+    **Assistant Token Masking Support:**
+
+    The class supports assistant token masking across multiple model families, allowing you to identify which tokens
+    in a conversation were generated by the assistant. This is crucial for reinforcement learning applications.
+
+    **Supported Model Families:**
+
+    - **Qwen family** (e.g., `Qwen/Qwen2.5-0.5B`): Custom template with full tool calling support
+    - **DialoGPT family** (e.g., `microsoft/DialoGPT-medium`): Custom template for conversation format
+    - **Falcon family** (e.g., `tiiuae/falcon-7b-instruct`): Custom template for instruction format
+    - **DeepSeek family** (e.g., `deepseek-ai/deepseek-coder-6.7b-base`): Custom template with native format
+    - **Other models** (OPT, GPT, MPT, BLOOM, Pythia, Phi, etc.): Default `chatml_format` template
+
+    **Example with Assistant Token Masking:**
+
+    .. code-block:: python
+
+        >>> from torchrl.data.llm.chat import History
+        >>> from torchrl.modules.llm.policies import ChatHistory
+        >>> from transformers import AutoTokenizer
+        >>>
+        >>> # Create a conversation history
+        >>> history = History.from_chats([[
+        ...     {"role": "user", "content": "Hello"},
+        ...     {"role": "assistant", "content": "Hi there!"},
+        ...     {"role": "user", "content": "How are you?"},
+        ...     {"role": "assistant", "content": "I'm doing well, thanks!"}
+        ... ]])
+        >>>
+        >>> # Create ChatHistory container for LLM wrapper
+        >>> chat_history = ChatHistory(prompt=history)
+        >>>
+        >>> # Load any supported tokenizer
+        >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")
+        >>>
+        >>> # Apply chat template with assistant token masking
+        >>> result = history.apply_chat_template(
+        ...     tokenizer=tokenizer,
+        ...     add_generation_prompt=False,
+        ...     return_dict=True,
+        ...     return_assistant_tokens_mask=True,
+        ... )
+        >>>
+        >>> # The result contains an assistant_masks tensor
+        >>> assistant_masks = result["assistant_masks"]
+        >>> print(f"Assistant tokens: {assistant_masks.sum().item()}")
+
+    **Integration with LLM Wrappers:**
+
+    History objects work seamlessly with the new modular wrapper design:
+
+    .. code-block:: python
+
+        >>> from torchrl.modules.llm import TransformersWrapper
+        >>> from torchrl.modules.llm.policies import ChatHistory
+        >>>
+        >>> # Create wrapper with history input mode
+        >>> wrapper = TransformersWrapper(
+        ...     model, tokenizer=tokenizer,
+        ...     input_mode="history",
+        ...     generate=True,
+        ...     return_log_probs=True
+        ... )
+        >>>
+        >>> # Use History with ChatHistory container
+        >>> history = History.from_chats([[
+        ...     {"role": "user", "content": "Hello"},
+        ...     {"role": "assistant", "content": "Hi there!"}
+        ... ]])
+        >>> chat_history = ChatHistory(prompt=history)
+        >>> result = wrapper(TensorDict(history=chat_history, batch_size=(1,)))
+        >>> print(result["history"].response)  # New response from LLM
+
+    Attributes:
+        role (str): The role of the message sender.
+        content (str): The content of the message.
+        is_complete (bool): Whether the message was properly terminated with an end token. Defaults to `True`.
+        tool_calls (list[dict] | None): Optional list of tool calls in the message.
+        tool_responses (list[str] | None): Optional list of tool responses.
+
+    Methods:
+        apply_chat_template: converts the `History` object to str / tokens.
+        append: append one element to the list of items along a given dimension.
+        extend: extend the list of items along a given dimension.
+
+    Examples:
+        >>> # With tensordict < 0.10, we need to tell the lib that lists constitute batches
+        >>> import tensordict
+        >>> tensordict.set_list_to_stack(True).set()
+        >>> import transformers
+        >>> history0 = History(
+        ...     role='system',
+        ...     content='''CONTENT
+        ... This is the setup''',
+        ... )
+        >>> history1 = History(
+        ...     role='user',
+        ...     content='''CONTENT
+        ... This is the first user prompt''',
+        ... )
+        >>> history2 = History(
+        ...     role='assistant',
+        ...     content='''CONTENT
+        ... This is the second prompt, the first for the assistant.''',
+        ... )
+        >>> history = torch.stack([history0, history1, history2])
+        >>> assert history.role == ['system', 'user', 'assistant']
+        >>> tokenizer = transformers.AutoTokenizer.from_pretrained("GPT2")
+        >>> # Apply a template to pass the history to an LLM. Note that the output has
+        >>> #  an additional prompt to elict an answer from the LLM thanks to the 'add_generation_prompt' argument.
+        >>> parsed_string = history.apply_chat_template(tokenizer=tokenizer, add_generation_prompt=True)
+        >>> parsed_string
+            <|im_start|>system
+        CONTENT
+        This is the setup<|im_end|>
+
+            <|im_start|>user
+        CONTENT
+        This is the first user prompt<|im_end|>
+
+            <|im_start|>assistant
+        CONTENT
+        This is the second prompt, the first for the assistant.<|im_end|>
+
+        <|im_start|>assistant
+
+    .. seealso::
+        :class:`~torchrl.modules.llm.policies.ChatHistory`: Container for managing conversation data in LLM environments.
+        :class:`~torchrl.modules.llm.policies.Text`: Container for text data.
+        :class:`~torchrl.modules.llm.policies.Tokens`: Container for token data.
+    """
+
+    role: str | list[str] | list[list[str]]
+    content: str | ContentBase | list[str] | list[ContentBase] | list[list[str]] | list[
+        list[ContentBase]
+    ]
+    is_complete: bool = True
+    tool_calls: list[dict] | None = None
+    tool_responses: list[str] | None = None
+
+    def __post_init__(self):
+        if not list_to_stack():
+            raise RuntimeError(
+                "Please set the list_to_stack to True using tensordict.set_list_to_stack(True).set() at the beginning of your script, "
+                "or the LIST_TO_STACK=1 environment variable."
+            )
+
+    def apply_chat_template(
+        self,
+        *,
+        tokenizer: transformers.AutoTokenizer | transformers.AutoProcessor,  # noqa
+        add_generation_prompt: bool = True,
+        chat_template: str | None = None,
+        chat_template_name: str | None = None,
+        continue_final_message: bool = False,
+        tokenize: bool | None = None,
+        padding: bool | str = False,
+        truncation: bool | str = False,
+        return_tensors: str | None = None,
+        return_dict: bool | None = None,
+        return_assistant_tokens_mask: bool = False,
+        **kwargs,
+    ) -> str | list[str] | TensorDict:
+        """Applies a chat template to the history.
+
+        Keyword Args:
+            tokenizer (transformers.PreTrainedTokenizer | transformers.AutoProcessor): The tokenizer to use.
+            add_generation_prompt (bool, optional): Whether to add a generation prompt (e.g. `"<|im_start|>assistant"`). Defaults to `True`.
+            chat_template (str, optional): The chat template to use. Defaults to the tokenizer's default template.
+            chat_template_name (str, optional): The name of the chat template to use.
+                Prevalent over `tokenizer.chat_template`. If `None`, the method will automatically detect the model family and use the appropriate template.
+                Defaults to `None`.
+            continue_final_message (bool, optional): Whether to continue the final message. Defaults to `False`.
+            tokenize (bool, optional): Whether to tokenize the output. Defaults to `False`.
+            padding (bool | str, optional): The padding strategy to use. Defaults to `False`.
+            truncation (bool | str, optional): The truncation strategy to use. Defaults to `False`.
+            return_tensors (str | None, optional): The type of tensors to return. Defaults to "pt".
+            return_dict (bool, optional): Whether to return a dictionary. Defaults to `False`.
+            return_assistant_tokens_mask (bool, optional): Whether to return a mask of the assistant generated tokens.
+                If `True`, the mask will be written to the `assistant_masks` key.
+                For tokens generated by the assistant, the mask will contain `1`.
+                For user and system tokens, the mask will contain `0`.
+                This functionality is only available for chat templates that support it via the `{% generation %}` keyword.
+                Defaults to `False`.
+
+                .. note:: Assistant token masking is supported across multiple model families:
+                    - **Qwen family**: Uses custom template with full tool calling support
+                    - **DialoGPT family**: Uses custom template for conversation format
+                    - **Falcon family**: Uses custom template for instruction format
+                    - **DeepSeek family**: Uses custom template with native format
+                    - **Other models**: Use the default `chatml_format` template
+
+                    The method automatically detects the model family and selects the appropriate template.
+
+            **kwargs: Additional keyword arguments to pass to the tokenizer `apply_chat_template` method.
+
+        Returns:
+            The formatted history.
+        """
+        if chat_template is None:
+            if chat_template_name is not None:
+                chat_template = _CHAT_TEMPLATES[chat_template_name]
+                chat_template_name = None
+            elif tokenizer is None:
+                raise RuntimeError(
+                    "You must specify a tokenizer to use when chat_template is not specified."
+                )
+            else:
+                # Auto-detect model family and use appropriate template
+                model_name = getattr(tokenizer, "name_or_path", "").lower()
+
+                # First check for custom model family keywords
+                custom_template_found = False
+                for template_name, keywords in _CUSTOM_MODEL_FAMILY_KEYWORDS.items():
+                    if any(keyword.lower() in model_name for keyword in keywords):
+                        chat_template = _CHAT_TEMPLATES[template_name]
+                        chat_template_name = None
+                        custom_template_found = True
+                        break
+
+                if not custom_template_found:
+                    # Fall back to built-in model family detection
+                    if "qwen" in model_name:
+                        # We prefer our implementation of the Qwen template,
+                        #  since it accounts for the assistant's masking.
+                        chat_template = _CHAT_TEMPLATES["qwen"]
+                        chat_template_name = None
+                    elif "dialogpt" in model_name or "microsoft/dialo" in model_name:
+                        # DialoGPT family - use our custom template
+                        chat_template = _CHAT_TEMPLATES["dialogpt"]
+                        chat_template_name = None
+                    elif "falcon" in model_name or "tiiuae/falcon" in model_name:
+                        # Falcon family - use our custom template
+                        chat_template = _CHAT_TEMPLATES["falcon"]
+                        chat_template_name = None
+                    elif "deepseek" in model_name:
+                        # DeepSeek family - use our custom template with generation keyword
+                        chat_template = _CHAT_TEMPLATES["deepseek"]
+                        chat_template_name = None
+                    elif "llama" in model_name:
+                        # Llama family - use our custom template
+                        chat_template = _CHAT_TEMPLATES["llama"]
+                        chat_template_name = None
+                    else:
+                        # For other models, check if their default template supports generation
+                        if (
+                            hasattr(tokenizer, "chat_template")
+                            and tokenizer.chat_template
+                            and "{% generation %}" in tokenizer.chat_template
+                        ):
+                            # Use the model's own template if it supports generation
+                            chat_template = tokenizer.chat_template
+                        else:
+                            # Use our default chatml_format template
+                            chat_template = _CHAT_TEMPLATES["chatml_format"]
+        if chat_template is None:
+            chat_template = _CHAT_TEMPLATES["chatml_format"]
+        if tokenize is None:
+            if return_assistant_tokens_mask or return_tensors is not None:
+                tokenize = True
+            else:
+                tokenize = False
+        if tokenize:
+            if return_tensors is None:
+                return_tensors = "pt"
+            if return_dict is None and return_assistant_tokens_mask:
+                return_dict = True
+        elif return_dict is None:
+            return_dict = False
+
+        if self.ndim > 1:
+            result = [
+                self[i].apply_chat_template(
+                    tokenizer=tokenizer,
+                    add_generation_prompt=add_generation_prompt,
+                    chat_template=chat_template,
+                    chat_template_name=chat_template_name,
+                    tokenize=tokenize,
+                    padding=padding,
+                    truncation=truncation,
+                    return_tensors=return_tensors,
+                    continue_final_message=continue_final_message,
+                    return_dict=return_dict,
+                    return_assistant_tokens_mask=return_assistant_tokens_mask,
+                    **kwargs,
+                )
+                for i in range(self.batch_size[0])
+            ]
+            if return_dict:
+                return lazy_stack(result)
+            else:
+                return result
+        self_flat = self.view(-1)
+        # tolist_first=True is needed to avoid having a list of dict of dicts, but a list of dicts of lists of dicts
+        self_flat = self_flat.tolist(tolist_first=True)
+        # Remove the "<none>" role
+        self_flat = [item for item in self_flat if item["role"] != "<none>"]
+        result = tokenizer.apply_chat_template(
+            conversation=self_flat,
+            add_generation_prompt=add_generation_prompt,
+            chat_template=chat_template,
+            tokenize=tokenize,
+            padding=padding,
+            truncation=truncation,
+            return_tensors=return_tensors,
+            continue_final_message=continue_final_message,
+            return_dict=return_dict,
+            return_assistant_tokens_mask=return_assistant_tokens_mask,
+            **kwargs,
+        )
+        if not isinstance(result, (torch.Tensor, list, str)):
+            result = TensorDict.from_dict(result, auto_batch_size=True, batch_dims=1)
+            # If self has a batch_dims of 1, we have just the time dimension, so we need to remove the batch dim from the result
+            if self.batch_dims == 1:
+                if result.batch_size[0] != 1:
+                    raise RuntimeError(
+                        f"Expected a batch size of 1, got {result.batch_size[0]}."
+                    )
+                result = result.squeeze(0)
+        return result
+
+    @classmethod
+    def from_text(
+        cls,
+        text: str | list[str],
+        chat_template_name: str | None = None,
+        # currently without effect
+        chat_template: str | None = None,
+        tokenizer: transformers.AutoTokenizer  # noqa: F821
+        | transformers.AutoProcessor  # noqa: F821
+        | None = None,
+    ) -> History:
+        r"""Inverts a chat template into a History object.
+
+        Args:
+            text (str | list[str]): The chat template to invert.
+            chat_template_name (str, optional): The name of the chat template to use.
+            tokenizer (transformers.AutoTokenizer | transformers.AutoProcessor, optional): The tokenizer to use.
+
+        Returns:
+            History: The inverted History object.
+
+        Examples:
+            >>> from torchrl.data.llm.history import History
+            >>> from transformers import AutoTokenizer
+            >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
+            >>> text = "<|im_start|>system\nYou are a helpful assistant.\n<|im_end|>\n<|im_start|>user\nWrite a python script that gives the capital of France or Germany.\n<|im_end|>\n<|im_start|>assistant\n<think>The capital of France is Paris, the capital of Germany is Berlin.</think>\n<answer><python>\n"
+            >>> history = History.from_text(text, tokenizer=tokenizer)
+            >>> print(history)
+            History(
+                content=NonTensorStack(
+                    ['You are a helpful assistant.', 'Write a python s...,
+                    batch_size=torch.Size([3]),
+                    device=None),
+                is_complete=NonTensorStack(
+                    [True, True, False],
+                    batch_size=torch.Size([3]),
+                    device=None),
+                role=NonTensorStack(
+                    ['system', 'user', 'assistant'],
+                    batch_size=torch.Size([3]),
+                    device=None),
+                tool_calls=None,
+                tool_responses=None,
+                batch_size=torch.Size([3]),
+                device=None,
+                is_shared=False)
+        """
+        if chat_template_name is None:
+            if chat_template is not None:
+                # TODO: find best match given template
+                pass
+
+            model_name = getattr(tokenizer, "name_or_path", "").lower()
+            # First check for custom model family keywords
+            custom_template_found = False
+            for template_name, keywords in _CUSTOM_MODEL_FAMILY_KEYWORDS.items():
+                if any(keyword.lower() in model_name for keyword in keywords):
+                    chat_template_name = template_name
+                    custom_template_found = True
+                    break
+
+            if not custom_template_found:
+                # Fall back to built-in model family detection
+                if "qwen" in model_name:
+                    # We can automatically detect the template name from the tokenizer
+                    #  and use the precoded parser.
+                    chat_template_name = "qwen"
+                elif "dialogpt" in model_name or "microsoft/dialo" in model_name:
+                    chat_template_name = "dialogpt"
+                elif "falcon" in model_name or "tiiuae/falcon" in model_name:
+                    chat_template_name = "falcon"
+                elif "deepseek" in model_name:
+                    chat_template_name = "deepseek"
+                elif "llama" in model_name:
+                    chat_template_name = "llama"
+                else:
+                    chat_template_name = "chatml_format"
+
+        # Get the appropriate inverse parser function
+        if chat_template_name in ("chatml_format",):
+            func = cls._inv_chatml
+        elif chat_template_name in ("qwen",):
+            func = cls._inv_qwen
+        elif chat_template_name in ("dialogpt",):
+            func = cls._inv_dialogpt
+        elif chat_template_name in ("falcon",):
+            func = cls._inv_falcon
+        elif chat_template_name in ("deepseek",):
+            func = cls._inv_deepseek
+        elif chat_template_name in ("llama",):
+            func = cls._inv_llama
+        elif chat_template_name in _CUSTOM_INVERSE_PARSERS:
+            # Use custom inverse parser
+            func = _CUSTOM_INVERSE_PARSERS[chat_template_name]
+        else:
+            raise NotImplementedError(
+                f"chat_template_name '{chat_template_name}' is not supported. "
+                "Supported templates: 'chatml_format', 'qwen', 'dialogpt', 'falcon', 'deepseek'. "
+                "Use add_chat_template() to add custom templates."
+            )
+        if isinstance(text, list):
+            list_of_histories = [func(t) for t in text]
+            try:
+                return lazy_stack(list_of_histories)
+            except RuntimeError as e:
+                raise RuntimeError(
+                    f"Failed to stack histories: {list_of_histories=}"
+                ) from e
+        return func(text)
+
+    @classmethod
+    def _inv_chatml(cls, text: str) -> History:
+        """Inverts a chatml string into a History object.
+
+        Args:
+            text (str): The chatml string to invert.
+
+        Returns:
+            History: The inverted History object.
+        """
+        import json
+
+        torchrl_logger.debug(f"Inverting chatml:\n{text}")
+        # Find all complete blocks (ending with im_end or endoftext)
+        complete_pattern = r"<\|im_start\|>(.*?)\n(.*?)<\|(im_end|endoftext)\|>"
+        complete_matches = re.findall(complete_pattern, text, flags=re.DOTALL)
+
+        # Find any incomplete block at the end
+        incomplete_pattern = r"<\|im_start\|>(.*?)\n(.*?)$"
+        incomplete_matches = []
+        if complete_matches:
+            # Look for incomplete block after the last complete one
+            last_complete = complete_matches[-1]
+            last_complete_text = f"<|im_start|>{last_complete[0]}\n{last_complete[1]}<|{last_complete[2]}|>"
+            remaining_text = text[
+                text.rindex(last_complete_text) + len(last_complete_text) :
+            ]
+            if remaining_text.strip():
+                incomplete_match = re.search(
+                    incomplete_pattern, remaining_text, flags=re.DOTALL
+                )
+                if incomplete_match:
+                    incomplete_matches = [
+                        (incomplete_match.group(1), incomplete_match.group(2), None)
+                    ]
+        else:
+            # No complete blocks, check entire text for incomplete block
+            incomplete_match = re.search(incomplete_pattern, text, flags=re.DOTALL)
+            if incomplete_match:
+                incomplete_matches = [
+                    (incomplete_match.group(1), incomplete_match.group(2), None)
+                ]
+
+        # Combine complete and incomplete matches
+        matches = complete_matches + incomplete_matches
+
+        # Define tool patterns - same as Qwen for consistency
+        tool_call_pattern = re.compile(r"<tool_call>\n(.*?)\n</tool_call>", re.DOTALL)
+        tool_response_pattern = re.compile(
+            r"<tool_response>\n(.*?)\n</tool_response>", re.DOTALL
+        )
+
+        parsed_messages = []
+        for match in matches:
+            role = match[0].strip()
+            content = match[1].strip()
+            is_complete = match[2] is not None  # None indicates incomplete
+
+            # Initialize message dict
+            message_dict = {
+                "role": role,
+                "content": content,
+                "is_complete": is_complete,
+                "tool_calls": None,
+                "tool_responses": None,
+            }
+
+            # Find tool calls within the message
+            tool_calls = tool_call_pattern.findall(content)
+            if tool_calls:
+                tool_calls_list = []
+                for tool_call in tool_calls:
+                    try:
+                        tool_call_dict = json.loads(tool_call)
+                        tool_calls_list.append(tool_call_dict)
+                    except json.JSONDecodeError:
+                        continue
+                if tool_calls_list:
+                    message_dict["tool_calls"] = tool_calls_list
+
+            # Check for tool responses
+            tool_responses = tool_response_pattern.findall(content)
+            if tool_responses:
+                message_dict["tool_responses"] = tool_responses
+
+            parsed_messages.append(cls(**message_dict))
+
+        if not parsed_messages:
+            raise RuntimeError(
+                f"Couldn't get a single item out of text {text}. A common cause "
+                f"if that special tokens should not be ommitted, did you set include_stop_str_in_output/skip_special_tokens=False?"
+            )
+
+        return lazy_stack(parsed_messages)
+
+    @classmethod
+    def _inv_qwen(cls, template):
+        import json
+
+        # Define regex patterns for different parts of the template
+        message_pattern = re.compile(
+            r"<\|im_start\|>(.*?)(?:<\|(im_end|endoftext)\|>|$)", re.DOTALL
+        )
+        tool_call_pattern = re.compile(r"<tool_call>\n(.*?)\n</tool_call>", re.DOTALL)
+        tool_response_pattern = re.compile(
+            r"<tool_response>\n(.*?)\n</tool_response>", re.DOTALL
+        )
+
+        # Find all messages and track if they end with a proper token
+        messages = []
+        is_complete_list = []
+        for match in message_pattern.finditer(template):
+            full_match = match.group(0)
+            messages.append(match.group(1))
+            # Check if the message ends with a proper token
+            is_complete_list.append(
+                full_match.endswith("<|im_end|>")
+                or full_match.endswith("<|endoftext|>")
+            )
+
+        parsed_messages = []
+        for message, is_complete in zip(messages, is_complete_list):
+            # Split the message into role and content
+            parts = message.split("\n", 1)
+            if len(parts) < 2:
+                continue
+            role, content = parts[0], parts[1]
+
+            # Initialize message dict
+            message_dict = {
+                "role": role.strip(),
+                "content": content.strip(),
+                "is_complete": is_complete,
+                "tool_calls": None,
+                "tool_responses": None,
+            }
+
+            # Find tool calls within the message
+            tool_calls = tool_call_pattern.findall(content)
+            if tool_calls:
+                tool_calls_list = []
+                for tool_call in tool_calls:
+                    try:
+                        tool_call_dict = json.loads(tool_call)
+                        tool_calls_list.append(tool_call_dict)
+                    except json.JSONDecodeError:
+                        continue
+                if tool_calls_list:
+                    message_dict["tool_calls"] = tool_calls_list
+
+            # Check for tool responses
+            tool_responses = tool_response_pattern.findall(content)
+            if tool_responses:
+                message_dict["tool_responses"] = tool_responses
+
+            parsed_messages.append(cls(**message_dict))
+
+        if not parsed_messages:
+            raise RuntimeError(
+                f"Couldn't get a single item out of text {template}. A common cause "
+                f"if that special tokens should not be ommitted, did you set include_stop_str_in_output/skip_special_tokens=False?"
+            )
+
+        return lazy_stack(parsed_messages)
+
+    @classmethod
+    def _inv_dialogpt(cls, text: str) -> History:
+        """Inverts a DialogPT string into a History object.
+
+        Args:
+            text (str): The DialogPT string to invert.
+
+        Returns:
+            History: The inverted History object.
+        """
+        torchrl_logger.debug(f"Inverting DialogPT:\n{text}")
+
+        # DialogPT format is simple: alternating user/assistant messages
+        # Split by lines and parse
+        lines = text.strip().split("\n")
+        parsed_messages = []
+
+        for line in lines:
+            line = line.strip()
+            if not line:
+                continue
+
+            # Determine role based on content
+            if line.startswith("Assistant:"):
+                role = "assistant"
+                content = line[len("Assistant:") :].strip()
+            elif line.startswith("User:"):
+                role = "user"
+                content = line[len("User:") :].strip()
+            else:
+                # Default to user if no prefix
+                role = "user"
+                content = line
+
+            message_dict = {
+                "role": role,
+                "content": content,
+                "is_complete": True,  # DialogPT doesn't have explicit end tokens
+                "tool_calls": None,
+                "tool_responses": None,
+            }
+
+            parsed_messages.append(cls(**message_dict))
+
+        if not parsed_messages:
+            raise RuntimeError(f"Couldn't get a single item out of text {text}.")
+
+        return lazy_stack(parsed_messages)
+
+    @classmethod
+    def _inv_falcon(cls, text: str) -> History:
+        """Inverts a Falcon string into a History object.
+
+        Args:
+            text (str): The Falcon string to invert.
+
+        Returns:
+            History: The inverted History object.
+        """
+        torchrl_logger.debug(f"Inverting Falcon:\n{text}")
+
+        # Falcon format: "User: ... Assistant: ..."
+        # Split by "User:" and "Assistant:" prefixes
+        import re
+
+        # Pattern to match User: and Assistant: messages
+        pattern = r"(User:|Assistant:)\s*(.*?)(?=(User:|Assistant:)|$)"
+        matches = re.findall(pattern, text, re.DOTALL)
+
+        parsed_messages = []
+        for match in matches:
+            if len(match) != 2:
+                continue
+            prefix, content = match
+            content = content.strip()
+            if not content:
+                continue
+
+            if prefix == "User:":
+                role = "user"
+            elif prefix == "Assistant:":
+                role = "assistant"
+            else:
+                continue
+
+            message_dict = {
+                "role": role,
+                "content": content,
+                "is_complete": True,  # Falcon doesn't have explicit end tokens
+                "tool_calls": None,
+                "tool_responses": None,
+            }
+
+            parsed_messages.append(cls(**message_dict))
+
+        if not parsed_messages:
+            raise RuntimeError(f"Couldn't get a single item out of text {text}.")
+
+        return lazy_stack(parsed_messages)
+
+    @classmethod
+    def _inv_deepseek(cls, text: str) -> History:
+        """Inverts a DeepSeek string into a History object.
+
+        Args:
+            text (str): The DeepSeek string to invert.
+
+        Returns:
+            History: The inverted History object.
+        """
+        torchrl_logger.debug(f"Inverting DeepSeek:\n{text}")
+        import re
+
+        # Remove leading/trailing special tokens (e.g.
+        text = re.sub(r"^<[^>]+>", "", text)  # Remove leading <...>
+        text = re.sub(r"<[^>]+>$", "", text)  # Remove trailing <...>
+        # Remove any REDACTED_SPECIAL_TOKEN if present
+        text = re.sub(r"REDACTED_SPECIAL_TOKEN", "", text)
+        # Pattern to match User: and Assistant: messages
+        pattern = r"(User:|Assistant:)\s*(.*?)(?=(User:|Assistant:)|$)"
+        matches = re.findall(pattern, text, re.DOTALL)
+        parsed_messages = []
+        for match in matches:
+            if len(match) < 2:
+                continue
+            prefix, content = match[0], match[1]
+            content = content.strip()
+            if not content:
+                continue
+            if prefix == "User:":
+                role = "user"
+            elif prefix == "Assistant:":
+                role = "assistant"
+            else:
+                continue
+            message_dict = {
+                "role": role,
+                "content": content,
+                "is_complete": True,  # DeepSeek doesn't have explicit end tokens
+                "tool_calls": None,
+                "tool_responses": None,
+            }
+            parsed_messages.append(cls(**message_dict))
+        if not parsed_messages:
+            raise RuntimeError(f"Couldn't get a single item out of text {text}.")
+        return lazy_stack(parsed_messages)
+
+    @classmethod
+    def _inv_llama(cls, text: str) -> History:
+        import re
+
+        messages = []
+
+        # Remove BOS token if present
+        if text.startswith("<|begin_of_text|>"):
+            text = text[len("<|begin_of_text|>") :]
+
+        # Pattern to match complete message blocks: <|header_start|>role<|header_end|>\n\ncontent<|eot|>
+        complete_pattern = r"<\|header_start\|>(\w+)<\|header_end\|>\n\n(.*?)<\|eot\|>"
+        complete_matches = re.findall(complete_pattern, text, re.DOTALL)
+
+        # Pattern to match incomplete message blocks: <|header_start|>role<|header_end|>\n\ncontent (without <|eot|>)
+        incomplete_pattern = r"<\|header_start\|>(\w+)<\|header_end\|>\n\n(.*?)$"
+
+        # Find any incomplete message at the end
+        incomplete_matches = []
+        if complete_matches:
+            # Look for incomplete message after the last complete one
+            last_complete_end = text.rfind("<|eot|>")
+            if last_complete_end != -1:
+                remaining_text = text[last_complete_end + len("<|eot|>") :]
+                if remaining_text.strip():
+                    incomplete_match = re.search(
+                        incomplete_pattern, remaining_text, re.DOTALL
+                    )
+                    if incomplete_match:
+                        incomplete_matches = [
+                            (
+                                incomplete_match.group(1),
+                                incomplete_match.group(2),
+                                False,
+                            )
+                        ]
+        else:
+            # No complete messages, check entire text for incomplete message
+            incomplete_match = re.search(incomplete_pattern, text, re.DOTALL)
+            if incomplete_match:
+                incomplete_matches = [
+                    (incomplete_match.group(1), incomplete_match.group(2), False)
+                ]
+
+        # Process complete messages
+        for role, content in complete_matches:
+            if content.strip():
+                messages.append(
+                    cls(role=role, content=content.strip(), is_complete=True)
+                )
+
+        # Process incomplete messages
+        for role, content, is_complete in incomplete_matches:
+            if content.strip():
+                messages.append(
+                    cls(role=role, content=content.strip(), is_complete=is_complete)
+                )
+
+        if not messages:
+            raise RuntimeError(f"Couldn't parse Llama format from text: {text}")
+
+        from tensordict import lazy_stack
+
+        return lazy_stack(messages)
+
+    def append(
+        self, history: History, *, inplace: bool = True, dim: int = -1
+    ) -> History:
+        """Appends a new history to the current one.
+
+        Args:
+            history (History): The new history to append.
+            inplace (bool, optional): Whether to perform the operation in-place. Defaults to `True`.
+            dim (int, optional): The dimension to append along. Defaults to -1.
+
+        Returns:
+            History: The appended History object.
+        """
+        # TODO: we should remove the <none> role from the history before appending / extending
+        #  It works when keeping them, but it may lead to a lot of useless padding in between valid messages
+        if not self.batch_dims:
+            raise RuntimeError(
+                "Cannot append an element to a batchless History. Call unsqueeze(dim=0) first on self."
+            )
+        if self.batch_dims != history.batch_dims + 1:
+            raise RuntimeError(
+                f"The new history to append must have one less dimension than self. Got self.ndim={self.ndim} and history.ndim={history.ndim}."
+            )
+        dim = _maybe_correct_neg_dim(dim, self.batch_size)
+        if inplace:
+            if (
+                isinstance(self._tensordict, LazyStackedTensorDict)
+                and self._tensordict.stack_dim == dim
+            ):
+                td = history._tensordict
+                if td.device != self.device:
+                    if self.device is None:
+                        td = td.copy().clear_device_()
+                    else:
+                        td = td.to(self.device)
+                self._tensordict.append(td)
+                return self
+            else:
+                td = history._tensordict
+                if td.device != self.device:
+                    if self.device is None:
+                        td = td.copy().clear_device_()
+                    else:
+                        td = td.to(self.device)
+                td = lazy_stack(list(self._tensordict.unbind(dim)) + [td], dim=dim)
+                self.__dict__["_tensordict"] = td
+                return self
+        if history.device != self.device:
+            if self.device is None:
+                history = history.copy().clear_device_()
+            else:
+                history = history.to(self.device)
+        return lazy_stack(list(self.unbind(dim)) + [history], dim=dim)
+
+    def extend(
+        self, history: History, *, inplace: bool = True, dim: int = 0
+    ) -> History:
+        if not self.batch_dims:
+            raise RuntimeError(
+                "Cannot add an element to a batchless History. Call unsqueeze(dim=0) first on self."
+            )
+        if self.batch_dims != history.batch_dims:
+            raise RuntimeError(
+                f"The new history to extend must have as many dimensions as self. Got self.ndim={self.ndim} and history.ndim={self.ndim}."
+            )
+        dim = _maybe_correct_neg_dim(dim, self.batch_size)
+        # if self.ndim > 1 and dim >= self.ndim - 1:
+        #     # then we need to append each element independently
+        #     result = []
+        #     for hist, new_hist in zip(self.unbind(0), history.unbind(0)):
+        #         hist_c = hist.extend(new_hist, inplace=inplace, dim=dim - 1)
+        #         result.append(hist_c)
+        #     if inplace:
+        #         return self
+        #     return lazy_stack(result)
+        if inplace:
+            if (
+                isinstance(self._tensordict, LazyStackedTensorDict)
+                and self._tensordict.stack_dim == dim
+            ):
+                td = history._tensordict
+                if td.device != self.device:
+                    if self.device is None:
+                        td = td.copy().clear_device_()
+                    else:
+                        td = td.to(self.device)
+                self._tensordict.extend(td)
+                return self
+            else:
+                td = lazy_stack(
+                    list(self._tensordict.unbind(dim))
+                    + list(history._tensordict.unbind(dim)),
+                    dim=dim,
+                )
+                if td.device != self.device:
+                    if self.device is None:
+                        td = td.copy().clear_device_()
+                    else:
+                        td = td.to(self.device)
+                self.__dict__["_tensordict"] = td
+                return self
+        if history.device != self.device:
+            if self.device is None:
+                history = history.copy().clear_device_()
+            else:
+                history = history.to(self.device)
+        return torch.stack(list(self.unbind(dim)) + list(history.unbind(dim)), dim=dim)
+
+    @classmethod
+    def default_spec(cls, shape=(-1,)):
+        """A default spec to use in transforms / envs that return History objects.
+
+        Args:
+            shape (torch.Size, optional): The shape of the returned History spec. Defaults to `(-1)` (variable length
+                along the time dimension).
+
+        Example:
+            >>> import tensordict
+            >>> from torchrl.data import History
+            >>> tensordict.set_list_to_stack(True).set()
+            >>>
+            >>> history = History(role=["system", "user"], content=["a message", "another message"], batch_size=(2,))
+            >>> spec = history.default_spec()
+            >>> print(spec)
+            Composite(
+                role: NonTensor(
+                    shape=torch.Size([-1]),
+                    space=None,
+                    device=None,
+                    dtype=None,
+                    domain=None,
+                    example_data=foo),
+                content: NonTensor(
+                    shape=torch.Size([-1]),
+                    space=None,
+                    device=None,
+                    dtype=None,
+                    domain=None,
+                    example_data=foo),
+                device=None,
+                shape=torch.Size([-1]))
+            >>> print(spec.zero())
+            History(
+                content=NonTensorData(data=foo, batch_size=torch.Size([1]), device=None),
+                role=NonTensorData(data=foo, batch_size=torch.Size([1]), device=None),
+                batch_size=torch.Size([1]),
+                device=None,
+                is_shared=False)
+
+        """
+        from torchrl.data import Composite, NonTensor
+
+        def get_default_value(field):
+            if field.default is not dataclasses.MISSING:
+                return field.default
+            elif field.type in (str, "str"):
+                return "foo"
+            else:
+                return None
+
+        defaults = {
+            k: NonTensor(
+                example_data=get_default_value(cls.__dataclass_fields__[k]),
+                shape=shape,
+            )
+            for k in cls.__dataclass_fields__
+        }
+
+        return Composite(defaults, shape=shape[:-1], data_cls=cls)
+
+    @classmethod
+    def from_chats(cls, chats: list[list[dict]]) -> History:
+        """Create a History object from a list of chats.
+
+        Args:
+            chats (list[list[dict]]): A list of chats, where each chat is a list of dictionaries.
+        """
+        if isinstance(chats[0], dict):
+            return lazy_stack([cls(**chat) for chat in chats])
+        else:
+            return lazy_stack([cls.from_chats(chat) for chat in chats])