lerobot/tests/processor/test_tokenizer_processor.py

#!/usr/bin/env python

# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Tests for the TokenizerProcessorStep class.
"""

import tempfile
from unittest.mock import patch

import pytest
import torch

from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
from lerobot.lerobot_types import TransitionKey
from lerobot.processor import DataProcessorPipeline, TokenizerProcessorStep
from lerobot.processor.converters import create_transition, identity_transition
from lerobot.utils.constants import (
    ACTION,
    OBS_IMAGE,
    OBS_LANGUAGE,
    OBS_LANGUAGE_SUBTASK_ATTENTION_MASK,
    OBS_LANGUAGE_SUBTASK_TOKENS,
    OBS_STATE,
)
from tests.utils import skip_if_package_missing


class MockTokenizer:
    """Mock tokenizer for testing that mimics transformers tokenizer interface."""

    def __init__(self, vocab_size: int = 1000):
        self.vocab_size = vocab_size

    def __call__(
        self,
        text: str | list[str],
        max_length: int = 512,
        truncation: bool = True,
        padding: str = "max_length",
        padding_side: str = "right",
        return_tensors: str = "pt",
        **kwargs,
    ) -> dict[str, torch.Tensor]:
        """Mock tokenization that returns deterministic tokens based on text."""
        texts = [text] if isinstance(text, str) else text

        batch_size = len(texts)

        # Create mock input_ids and attention_mask
        input_ids = torch.zeros(batch_size, max_length, dtype=torch.long)
        attention_mask = torch.zeros(batch_size, max_length, dtype=torch.long)

        for i, txt in enumerate(texts):
            # Simple mock: use hash of text to generate deterministic tokens
            text_hash = hash(txt) % self.vocab_size
            seq_len = min(len(txt.split()), max_length)

            # Fill input_ids with simple pattern based on text
            for j in range(seq_len):
                input_ids[i, j] = (text_hash + j) % self.vocab_size

            # Set attention mask for non-padded positions
            attention_mask[i, :seq_len] = 1

        result = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
        }

        # Return single sequence for single input to match transformers behavior
        if len(texts) == 1:
            result = {k: v.squeeze(0) for k, v in result.items()}

        return result


@pytest.fixture
def mock_tokenizer():
    """Provide a mock tokenizer for testing."""
    return MockTokenizer(vocab_size=100)


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_basic_tokenization(mock_auto_tokenizer):
    """Test basic string tokenization functionality."""
    # Mock AutoTokenizer.from_pretrained to return our mock tokenizer
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=10)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "pick up the red cube"},
    )

    result = processor(transition)

    # Check that original task is preserved
    assert result[TransitionKey.COMPLEMENTARY_DATA]["task"] == "pick up the red cube"

    # Check that tokens were added to observation
    observation = result[TransitionKey.OBSERVATION]
    assert f"{OBS_LANGUAGE}.tokens" in observation
    assert f"{OBS_LANGUAGE}.attention_mask" in observation

    # Check token structure
    tokens = observation[f"{OBS_LANGUAGE}.tokens"]
    attention_mask = observation[f"{OBS_LANGUAGE}.attention_mask"]
    assert isinstance(tokens, torch.Tensor)
    assert isinstance(attention_mask, torch.Tensor)
    assert tokens.shape == (10,)
    assert attention_mask.shape == (10,)


@skip_if_package_missing("transformers")
def test_basic_tokenization_with_tokenizer_object():
    """Test basic string tokenization functionality using tokenizer object directly."""
    mock_tokenizer = MockTokenizer(vocab_size=100)

    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "pick up the red cube"},
    )

    result = processor(transition)

    # Check that original task is preserved
    assert result[TransitionKey.COMPLEMENTARY_DATA]["task"] == "pick up the red cube"

    # Check that tokens were added to observation
    observation = result[TransitionKey.OBSERVATION]
    assert f"{OBS_LANGUAGE}.tokens" in observation
    assert f"{OBS_LANGUAGE}.attention_mask" in observation

    # Check token structure
    tokens = observation[f"{OBS_LANGUAGE}.tokens"]
    attention_mask = observation[f"{OBS_LANGUAGE}.attention_mask"]
    assert isinstance(tokens, torch.Tensor)
    assert isinstance(attention_mask, torch.Tensor)
    assert tokens.shape == (10,)
    assert attention_mask.shape == (10,)


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_list_of_strings_tokenization(mock_auto_tokenizer):
    """Test tokenization of a list of strings."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=8)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": ["pick up cube", "place on table"]},
    )

    result = processor(transition)

    # Check that original task is preserved
    assert result[TransitionKey.COMPLEMENTARY_DATA]["task"] == ["pick up cube", "place on table"]

    # Check that tokens were added to observation
    observation = result[TransitionKey.OBSERVATION]
    tokens = observation[f"{OBS_LANGUAGE}.tokens"]
    attention_mask = observation[f"{OBS_LANGUAGE}.attention_mask"]
    assert tokens.shape == (2, 8)  # batch_size=2, seq_len=8
    assert attention_mask.shape == (2, 8)


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_tuple_of_strings_tokenization(mock_auto_tokenizer):
    """Test tokenization of a tuple of strings (returned by VectorEnv.call())."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=8)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": ("pick up cube", "place on table")},
    )

    result = processor(transition)

    observation = result[TransitionKey.OBSERVATION]
    tokens = observation[f"{OBS_LANGUAGE}.tokens"]
    attention_mask = observation[f"{OBS_LANGUAGE}.attention_mask"]
    assert tokens.shape == (2, 8)
    assert attention_mask.shape == (2, 8)


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_custom_keys(mock_auto_tokenizer):
    """Test using custom task_key."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", task_key="instruction", max_length=5)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"instruction": "move forward"},
    )

    result = processor(transition)

    # Check that tokens are stored in observation regardless of task_key
    observation = result[TransitionKey.OBSERVATION]
    assert f"{OBS_LANGUAGE}.tokens" in observation
    assert f"{OBS_LANGUAGE}.attention_mask" in observation

    tokens = observation[f"{OBS_LANGUAGE}.tokens"]
    assert tokens.shape == (5,)


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_none_complementary_data(mock_auto_tokenizer):
    """Test handling of None complementary_data."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer")

    transition = create_transition(observation={}, complementary_data=None)

    # create_transition converts None complementary_data to empty dict, so task key is missing
    with pytest.raises(KeyError, match="task"):
        processor(transition)


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_missing_task_key(mock_auto_tokenizer):
    """Test handling when task key is missing."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer")

    transition = create_transition(observation={}, complementary_data={"other_field": "some value"})

    with pytest.raises(KeyError, match="task"):
        processor(transition)


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_none_task_value(mock_auto_tokenizer):
    """Test handling when task value is None."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer")

    transition = create_transition(observation={}, complementary_data={"task": None})

    with pytest.raises(ValueError, match="Task extracted from Complementary data is None"):
        processor(transition)


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_unsupported_task_type(mock_auto_tokenizer):
    """Test handling of unsupported task types."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer")

    # Test with integer task - get_task returns None, observation raises ValueError
    transition = create_transition(observation={}, complementary_data={"task": 123})

    with pytest.raises(ValueError, match="Task cannot be None"):
        processor(transition)

    # Test with mixed list - get_task returns None, observation raises ValueError
    transition = create_transition(observation={}, complementary_data={"task": ["text", 123, "more text"]})

    with pytest.raises(ValueError, match="Task cannot be None"):
        processor(transition)


@skip_if_package_missing("transformers")
def test_no_tokenizer_error():
    """Test that ValueError is raised when neither tokenizer nor tokenizer_name is provided."""
    with pytest.raises(ValueError, match="Either 'tokenizer' or 'tokenizer_name' must be provided"):
        TokenizerProcessorStep()


@skip_if_package_missing("transformers")
def test_invalid_tokenizer_name_error():
    """Test that error is raised when invalid tokenizer_name is provided."""
    with patch("lerobot.processor.tokenizer_processor.AutoTokenizer") as mock_auto_tokenizer:
        # Mock import error
        mock_auto_tokenizer.from_pretrained.side_effect = Exception("Model not found")

        with pytest.raises(Exception, match="Model not found"):
            TokenizerProcessorStep(tokenizer_name="invalid-tokenizer")


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_get_config_with_tokenizer_name(mock_auto_tokenizer):
    """Test configuration serialization when using tokenizer_name."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    processor = TokenizerProcessorStep(
        tokenizer_name="test-tokenizer",
        max_length=256,
        task_key="instruction",
        padding="longest",
        truncation=False,
    )

    config = processor.get_config()

    expected = {
        "tokenizer_name": "test-tokenizer",
        "max_length": 256,
        "task_key": "instruction",
        "padding_side": "right",
        "padding": "longest",
        "truncation": False,
    }

    assert config == expected


@skip_if_package_missing("transformers")
def test_get_config_with_tokenizer_object():
    """Test configuration serialization when using tokenizer object."""
    mock_tokenizer = MockTokenizer(vocab_size=100)

    processor = TokenizerProcessorStep(
        tokenizer=mock_tokenizer,
        max_length=256,
        task_key="instruction",
        padding="longest",
        truncation=False,
    )

    config = processor.get_config()

    # tokenizer_name should not be in config when tokenizer object is used
    expected = {
        "max_length": 256,
        "task_key": "instruction",
        "padding_side": "right",
        "padding": "longest",
        "truncation": False,
    }

    assert config == expected
    assert "tokenizer_name" not in config


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_state_dict_methods(mock_auto_tokenizer):
    """Test state_dict and load_state_dict methods."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer")

    # Should return empty dict
    state = processor.state_dict()
    assert state == {}

    # load_state_dict should not raise error
    processor.load_state_dict({})


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_reset_method(mock_auto_tokenizer):
    """Test reset method."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer")

    # Should not raise error
    processor.reset()


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_integration_with_robot_processor(mock_auto_tokenizer):
    """Test integration with RobotProcessor."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    tokenizer_processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=6)
    robot_processor = DataProcessorPipeline(
        [tokenizer_processor], to_transition=identity_transition, to_output=identity_transition
    )

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "test task"},
    )

    result = robot_processor(transition)

    # Check that observation exists and tokenization was applied
    assert TransitionKey.OBSERVATION in result
    observation = result[TransitionKey.OBSERVATION]
    assert f"{OBS_LANGUAGE}.tokens" in observation
    assert f"{OBS_LANGUAGE}.attention_mask" in observation
    tokens = observation[f"{OBS_LANGUAGE}.tokens"]
    attention_mask = observation[f"{OBS_LANGUAGE}.attention_mask"]
    assert tokens.shape == (6,)
    assert attention_mask.shape == (6,)

    # Check that other data is preserved
    assert torch.equal(
        result[TransitionKey.OBSERVATION]["state"], transition[TransitionKey.OBSERVATION]["state"]
    )
    assert torch.equal(result[TransitionKey.ACTION], transition[TransitionKey.ACTION])


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_save_and_load_pretrained_with_tokenizer_name(mock_auto_tokenizer):
    """Test saving and loading processor with tokenizer_name."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    original_processor = TokenizerProcessorStep(
        tokenizer_name="test-tokenizer", max_length=32, task_key="instruction"
    )

    robot_processor = DataProcessorPipeline(
        [original_processor], to_transition=identity_transition, to_output=identity_transition
    )

    with tempfile.TemporaryDirectory() as temp_dir:
        # Save processor
        robot_processor.save_pretrained(temp_dir)

        # Load processor - tokenizer will be recreated from saved config
        loaded_processor = DataProcessorPipeline.from_pretrained(
            temp_dir,
            config_filename="dataprocessorpipeline.json",
            to_transition=identity_transition,
            to_output=identity_transition,
        )

        # Test that loaded processor works
        transition = create_transition(
            observation={"state": torch.tensor([1.0, 2.0])},
            action=torch.tensor([0.1, 0.2]),
            complementary_data={"instruction": "test instruction"},
        )

        result = loaded_processor(transition)
        assert TransitionKey.OBSERVATION in result
        assert f"{OBS_LANGUAGE}.tokens" in result[TransitionKey.OBSERVATION]
        assert f"{OBS_LANGUAGE}.attention_mask" in result[TransitionKey.OBSERVATION]


@skip_if_package_missing("transformers")
def test_save_and_load_pretrained_with_tokenizer_object():
    """Test saving and loading processor with tokenizer object using overrides."""
    mock_tokenizer = MockTokenizer(vocab_size=100)

    original_processor = TokenizerProcessorStep(
        tokenizer=mock_tokenizer, max_length=32, task_key="instruction"
    )

    robot_processor = DataProcessorPipeline(
        [original_processor], to_transition=identity_transition, to_output=identity_transition
    )

    with tempfile.TemporaryDirectory() as temp_dir:
        # Save processor
        robot_processor.save_pretrained(temp_dir)

        # Load processor with tokenizer override (since tokenizer object wasn't saved)
        loaded_processor = DataProcessorPipeline.from_pretrained(
            temp_dir,
            config_filename="dataprocessorpipeline.json",
            overrides={"tokenizer_processor": {"tokenizer": mock_tokenizer}},
            to_transition=identity_transition,
            to_output=identity_transition,
        )

        # Test that loaded processor works
        transition = create_transition(
            observation={"state": torch.tensor([1.0, 2.0])},
            action=torch.tensor([0.1, 0.2]),
            complementary_data={"instruction": "test instruction"},
        )

        result = loaded_processor(transition)
        assert TransitionKey.OBSERVATION in result
        assert f"{OBS_LANGUAGE}.tokens" in result[TransitionKey.OBSERVATION]
        assert f"{OBS_LANGUAGE}.attention_mask" in result[TransitionKey.OBSERVATION]


@skip_if_package_missing("transformers")
def test_registry_functionality():
    """Test that the processor is properly registered."""
    from lerobot.processor import ProcessorStepRegistry

    # Check that the processor is registered
    assert "tokenizer_processor" in ProcessorStepRegistry.list()

    # Check that we can retrieve it
    retrieved_class = ProcessorStepRegistry.get("tokenizer_processor")
    assert retrieved_class is TokenizerProcessorStep


@skip_if_package_missing("transformers")
def test_features_basic():
    """Test basic feature contract functionality."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=128)

    input_features = {
        PipelineFeatureType.OBSERVATION: {OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(10,))},
        PipelineFeatureType.ACTION: {ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(5,))},
    }

    output_features = processor.transform_features(input_features)

    # Check that original features are preserved
    assert OBS_STATE in output_features[PipelineFeatureType.OBSERVATION]
    assert ACTION in output_features[PipelineFeatureType.ACTION]

    # Check that tokenized features are added
    assert f"{OBS_LANGUAGE}.tokens" in output_features[PipelineFeatureType.OBSERVATION]
    assert f"{OBS_LANGUAGE}.attention_mask" in output_features[PipelineFeatureType.OBSERVATION]

    # Check feature properties
    tokens_feature = output_features[PipelineFeatureType.OBSERVATION][f"{OBS_LANGUAGE}.tokens"]
    attention_mask_feature = output_features[PipelineFeatureType.OBSERVATION][
        f"{OBS_LANGUAGE}.attention_mask"
    ]

    assert tokens_feature.type == FeatureType.LANGUAGE
    assert tokens_feature.shape == (128,)
    assert attention_mask_feature.type == FeatureType.LANGUAGE
    assert attention_mask_feature.shape == (128,)


@skip_if_package_missing("transformers")
def test_features_with_custom_max_length():
    """Test feature contract with custom max_length."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=64)

    input_features = {PipelineFeatureType.OBSERVATION: {}}
    output_features = processor.transform_features(input_features)

    # Check that features use correct max_length
    assert f"{OBS_LANGUAGE}.tokens" in output_features[PipelineFeatureType.OBSERVATION]
    assert f"{OBS_LANGUAGE}.attention_mask" in output_features[PipelineFeatureType.OBSERVATION]

    tokens_feature = output_features[PipelineFeatureType.OBSERVATION][f"{OBS_LANGUAGE}.tokens"]
    attention_mask_feature = output_features[PipelineFeatureType.OBSERVATION][
        f"{OBS_LANGUAGE}.attention_mask"
    ]

    assert tokens_feature.shape == (64,)
    assert attention_mask_feature.shape == (64,)


@skip_if_package_missing("transformers")
def test_features_existing_features():
    """Test feature contract when tokenized features already exist."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=256)

    input_features = {
        PipelineFeatureType.OBSERVATION: {
            f"{OBS_LANGUAGE}.tokens": PolicyFeature(type=FeatureType.LANGUAGE, shape=(100,)),
            f"{OBS_LANGUAGE}.attention_mask": PolicyFeature(type=FeatureType.LANGUAGE, shape=(100,)),
        }
    }

    output_features = processor.transform_features(input_features)

    # Should not overwrite existing features
    assert output_features[PipelineFeatureType.OBSERVATION][f"{OBS_LANGUAGE}.tokens"].shape == (
        100,
    )  # Original shape preserved
    assert output_features[PipelineFeatureType.OBSERVATION][f"{OBS_LANGUAGE}.attention_mask"].shape == (100,)


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_tokenization_parameters(mock_auto_tokenizer):
    """Test that tokenization parameters are correctly passed to tokenizer."""

    # Create a custom mock that tracks calls
    class TrackingMockTokenizer:
        def __init__(self):
            self.last_call_args = None
            self.last_call_kwargs = None

        def __call__(self, *args, **kwargs):
            self.last_call_args = args
            self.last_call_kwargs = kwargs
            # Return minimal valid output
            return {
                "input_ids": torch.zeros(16, dtype=torch.long),
                "attention_mask": torch.ones(16, dtype=torch.long),
            }

    tracking_tokenizer = TrackingMockTokenizer()
    mock_auto_tokenizer.from_pretrained.return_value = tracking_tokenizer

    processor = TokenizerProcessorStep(
        tokenizer_name="test-tokenizer",
        max_length=16,
        padding="longest",
        truncation=False,
        padding_side="left",
    )

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "test task"},
    )

    processor(transition)

    # Check that parameters were passed correctly (task is converted to list)
    assert tracking_tokenizer.last_call_args == (["test task"],)
    assert tracking_tokenizer.last_call_kwargs["max_length"] == 16
    assert tracking_tokenizer.last_call_kwargs["padding"] == "longest"
    assert tracking_tokenizer.last_call_kwargs["padding_side"] == "left"
    assert tracking_tokenizer.last_call_kwargs["truncation"] is False
    assert tracking_tokenizer.last_call_kwargs["return_tensors"] == "pt"


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_preserves_other_complementary_data(mock_auto_tokenizer):
    """Test that other complementary data fields are preserved."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer")

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={
            "task": "test task",
            "episode_id": 123,
            "timestamp": 456.789,
            "other_field": {"nested": "data"},
        },
    )

    result = processor(transition)
    comp_data = result[TransitionKey.COMPLEMENTARY_DATA]

    # Check that all original fields are preserved
    assert comp_data["task"] == "test task"
    assert comp_data["episode_id"] == 123
    assert comp_data["timestamp"] == 456.789
    assert comp_data["other_field"] == {"nested": "data"}

    # Check that tokens were added to observation
    observation = result[TransitionKey.OBSERVATION]
    assert f"{OBS_LANGUAGE}.tokens" in observation
    assert f"{OBS_LANGUAGE}.attention_mask" in observation


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_deterministic_tokenization(mock_auto_tokenizer):
    """Test that tokenization is deterministic for the same input."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=10)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "consistent test"},
    )

    result1 = processor(transition)
    result2 = processor(transition)

    tokens1 = result1[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.tokens"]
    attention_mask1 = result1[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.attention_mask"]
    tokens2 = result2[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.tokens"]
    attention_mask2 = result2[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.attention_mask"]

    # Results should be identical
    assert torch.equal(tokens1, tokens2)
    assert torch.equal(attention_mask1, attention_mask2)


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_empty_string_task(mock_auto_tokenizer):
    """Test handling of empty string task."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=8)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": ""},
    )

    result = processor(transition)

    # Should still tokenize (mock tokenizer handles empty strings)
    observation = result[TransitionKey.OBSERVATION]
    assert f"{OBS_LANGUAGE}.tokens" in observation
    tokens = observation[f"{OBS_LANGUAGE}.tokens"]
    assert tokens.shape == (8,)


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_very_long_task(mock_auto_tokenizer):
    """Test handling of very long task strings."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=5, truncation=True)

    long_task = " ".join(["word"] * 100)  # Very long task
    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": long_task},
    )

    result = processor(transition)

    # Should be truncated to max_length
    observation = result[TransitionKey.OBSERVATION]
    tokens = observation[f"{OBS_LANGUAGE}.tokens"]
    attention_mask = observation[f"{OBS_LANGUAGE}.attention_mask"]
    assert tokens.shape == (5,)
    assert attention_mask.shape == (5,)


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_custom_padding_side(mock_auto_tokenizer):
    """Test using custom padding_side parameter."""

    # Create a mock tokenizer that tracks padding_side calls
    class PaddingSideTrackingTokenizer:
        def __init__(self):
            self.padding_side_calls = []

        def __call__(
            self,
            text,
            max_length=512,
            truncation=True,
            padding="max_length",
            padding_side="right",
            return_tensors="pt",
            **kwargs,
        ):
            self.padding_side_calls.append(padding_side)
            # Return minimal valid output
            return {
                "input_ids": torch.zeros(max_length, dtype=torch.long),
                "attention_mask": torch.ones(max_length, dtype=torch.long),
            }

    tracking_tokenizer = PaddingSideTrackingTokenizer()
    mock_auto_tokenizer.from_pretrained.return_value = tracking_tokenizer

    # Test left padding
    processor_left = TokenizerProcessorStep(
        tokenizer_name="test-tokenizer", max_length=10, padding_side="left"
    )

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "test task"},
    )
    processor_left(transition)

    assert tracking_tokenizer.padding_side_calls[-1] == "left"

    # Test right padding (default)
    processor_right = TokenizerProcessorStep(
        tokenizer_name="test-tokenizer", max_length=10, padding_side="right"
    )

    processor_right(transition)

    assert tracking_tokenizer.padding_side_calls[-1] == "right"


@skip_if_package_missing("transformers")
def test_device_detection_cpu():
    """Test that tokenized tensors stay on CPU when other tensors are on CPU."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    # Create transition with CPU tensors
    observation = {OBS_STATE: torch.randn(10)}  # CPU tensor
    action = torch.randn(5)  # CPU tensor
    transition = create_transition(
        observation=observation, action=action, complementary_data={"task": "test task"}
    )

    result = processor(transition)

    # Check that tokenized tensors are on CPU
    tokens = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.tokens"]
    attention_mask = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.attention_mask"]

    assert tokens.device.type == "cpu"
    assert attention_mask.device.type == "cpu"


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
@skip_if_package_missing("transformers")
def test_device_detection_cuda():
    """Test that tokenized tensors are moved to CUDA when other tensors are on CUDA."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    # Create transition with CUDA tensors
    observation = {OBS_STATE: torch.randn(10).cuda()}  # CUDA tensor
    action = torch.randn(5).cuda()  # CUDA tensor
    transition = create_transition(
        observation=observation, action=action, complementary_data={"task": "test task"}
    )

    result = processor(transition)

    # Check that tokenized tensors are on CUDA
    tokens = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.tokens"]
    attention_mask = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.attention_mask"]

    assert tokens.device.type == "cuda"
    assert attention_mask.device.type == "cuda"
    assert tokens.device.index == 0  # Should be on same device as input


@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Requires at least 2 GPUs")
@skip_if_package_missing("transformers")
def test_device_detection_multi_gpu():
    """Test that tokenized tensors match device in multi-GPU setup."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    # Test with tensors on cuda:1
    device = torch.device("cuda:1")
    observation = {OBS_STATE: torch.randn(10).to(device)}
    action = torch.randn(5).to(device)
    transition = create_transition(
        observation=observation, action=action, complementary_data={"task": "multi gpu test"}
    )

    result = processor(transition)

    # Check that tokenized tensors are on cuda:1
    tokens = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.tokens"]
    attention_mask = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.attention_mask"]

    assert tokens.device == device
    assert attention_mask.device == device


@skip_if_package_missing("transformers")
def test_device_detection_no_tensors():
    """Test that tokenized tensors stay on CPU when no other tensors exist."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    # Create transition with no tensors
    transition = create_transition(
        observation={"metadata": {"key": "value"}},  # No tensors
        complementary_data={"task": "no tensor test"},
    )

    result = processor(transition)

    # Check that tokenized tensors are on CPU (default)
    tokens = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.tokens"]
    attention_mask = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.attention_mask"]

    assert tokens.device.type == "cpu"
    assert attention_mask.device.type == "cpu"


@skip_if_package_missing("transformers")
def test_device_detection_mixed_devices():
    """Test device detection when tensors are on different devices (uses first found)."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    if torch.cuda.is_available():
        # Create transition with mixed devices
        observation = {
            "observation.cpu": torch.randn(10),  # CPU
            "observation.cuda": torch.randn(10).cuda(),  # CUDA
        }
        transition = create_transition(
            observation=observation, complementary_data={"task": "mixed device test"}
        )

        result = processor(transition)

        # The device detection should use the first tensor found
        # (iteration order depends on dict, but result should be consistent)
        tokens = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.tokens"]
        attention_mask = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.attention_mask"]

        # Both should be on the same device
        assert tokens.device == attention_mask.device


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
@skip_if_package_missing("transformers")
def test_device_detection_from_action():
    """Test that device is detected from action tensor when no observation tensors exist."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    # Create transition with action on CUDA but no observation tensors
    observation = {"metadata": {"key": "value"}}  # No tensors in observation
    action = torch.randn(5).cuda()
    transition = create_transition(
        observation=observation, action=action, complementary_data={"task": "action device test"}
    )

    result = processor(transition)

    # Check that tokenized tensors match action's device
    tokens = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.tokens"]
    attention_mask = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.attention_mask"]

    assert tokens.device.type == "cuda"
    assert attention_mask.device.type == "cuda"


@skip_if_package_missing("transformers")
def test_device_detection_preserves_dtype():
    """Test that device detection doesn't affect dtype of tokenized tensors."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    # Create transition with float tensor (to test dtype isn't affected)
    observation = {OBS_STATE: torch.randn(10, dtype=torch.float16)}
    transition = create_transition(observation=observation, complementary_data={"task": "dtype test"})

    result = processor(transition)

    # Check that tokenized tensors have correct dtypes (not affected by input dtype)
    tokens = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.tokens"]
    attention_mask = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.attention_mask"]

    assert tokens.dtype == torch.long  # Should remain long
    assert attention_mask.dtype == torch.bool  # Should be bool (converted in processor)


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_integration_with_device_processor(mock_auto_tokenizer):
    """Test that TokenizerProcessorStep works correctly with DeviceProcessorStep in pipeline."""
    from lerobot.processor import DeviceProcessorStep

    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    # Create pipeline with TokenizerProcessorStep then DeviceProcessorStep
    tokenizer_processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=6)
    device_processor = DeviceProcessorStep(device="cuda:0")
    robot_processor = DataProcessorPipeline(
        [tokenizer_processor, device_processor],
        to_transition=identity_transition,
        to_output=identity_transition,
    )

    # Start with CPU tensors
    transition = create_transition(
        observation={OBS_STATE: torch.randn(10)},  # CPU
        action=torch.randn(5),  # CPU
        complementary_data={"task": "pipeline test"},
    )

    result = robot_processor(transition)

    # All tensors should end up on CUDA (moved by DeviceProcessorStep)
    assert result[TransitionKey.OBSERVATION][OBS_STATE].device.type == "cuda"
    assert result[TransitionKey.ACTION].device.type == "cuda"

    # Tokenized tensors should also be on CUDA
    tokens = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.tokens"]
    attention_mask = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.attention_mask"]
    assert tokens.device.type == "cuda"
    assert attention_mask.device.type == "cuda"


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
@skip_if_package_missing("transformers")
def test_simulated_accelerate_scenario():
    """Test scenario simulating Accelerate with data already on GPU."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    # Simulate Accelerate scenario: batch already on GPU
    device = torch.device("cuda:0")
    observation = {
        OBS_STATE: torch.randn(1, 10).to(device),  # Batched, on GPU
        OBS_IMAGE: torch.randn(1, 3, 224, 224).to(device),  # Batched, on GPU
    }
    action = torch.randn(1, 5).to(device)  # Batched, on GPU

    transition = create_transition(
        observation=observation,
        action=action,
        complementary_data={"task": ["accelerate test"]},  # List for batched task
    )

    result = processor(transition)

    # Tokenized tensors should match GPU placement
    tokens = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.tokens"]
    attention_mask = result[TransitionKey.OBSERVATION][f"{OBS_LANGUAGE}.attention_mask"]

    assert tokens.device == device
    assert attention_mask.device == device
    # MockTokenizer squeezes single-item batches, so shape is (max_length,) not (1, max_length)
    assert tokens.shape == (10,)  # MockTokenizer behavior for single string in list
    assert attention_mask.shape == (10,)


# =============================================================================
# Tests for get_subtask method
# =============================================================================


@skip_if_package_missing("transformers")
def test_get_subtask_missing_key():
    """Test get_subtask returns None when subtask key is missing from complementary_data."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task"},  # No "subtask" key
    )

    result = processor.get_subtask(transition)
    assert result is None


@skip_if_package_missing("transformers")
def test_get_subtask_none_value():
    """Test get_subtask returns None when subtask value is None."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task", "subtask": None},
    )

    result = processor.get_subtask(transition)
    assert result is None


@skip_if_package_missing("transformers")
def test_get_subtask_none_complementary_data():
    """Test get_subtask returns None when complementary_data is None."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data=None,  # No complementary data
    )

    result = processor.get_subtask(transition)
    assert result is None


@skip_if_package_missing("transformers")
def test_get_subtask_string():
    """Test get_subtask returns list with single string when subtask is a string."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task", "subtask": "pick up the cube"},
    )

    result = processor.get_subtask(transition)
    assert result == ["pick up the cube"]
    assert isinstance(result, list)
    assert len(result) == 1


@skip_if_package_missing("transformers")
def test_get_subtask_list_of_strings():
    """Test get_subtask returns the list when subtask is already a list of strings."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    subtask_list = ["pick up", "move to target", "place down"]
    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task", "subtask": subtask_list},
    )

    result = processor.get_subtask(transition)
    assert result == subtask_list
    assert isinstance(result, list)
    assert len(result) == 3


@skip_if_package_missing("transformers")
def test_get_subtask_unsupported_type_integer():
    """Test get_subtask returns None when subtask is an unsupported type (integer)."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task", "subtask": 123},
    )

    result = processor.get_subtask(transition)
    assert result is None


@skip_if_package_missing("transformers")
def test_get_subtask_unsupported_type_mixed_list():
    """Test get_subtask returns None when subtask is a list with mixed types."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task", "subtask": ["valid string", 123, "another string"]},
    )

    result = processor.get_subtask(transition)
    assert result is None


@skip_if_package_missing("transformers")
def test_get_subtask_unsupported_type_dict():
    """Test get_subtask returns None when subtask is a dictionary."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task", "subtask": {"key": "value"}},
    )

    result = processor.get_subtask(transition)
    assert result is None


@skip_if_package_missing("transformers")
def test_get_subtask_empty_string():
    """Test get_subtask with empty string returns list with empty string."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task", "subtask": ""},
    )

    result = processor.get_subtask(transition)
    assert result == [""]


@skip_if_package_missing("transformers")
def test_get_subtask_empty_list():
    """Test get_subtask with empty list returns empty list."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task", "subtask": []},
    )

    result = processor.get_subtask(transition)
    assert result == []


# =============================================================================
# Tests for subtask tokenization in observation method
# =============================================================================


@skip_if_package_missing("transformers")
def test_subtask_tokenization_when_present():
    """Test that subtask is tokenized and added to observation when present."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=8)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task", "subtask": "pick up the red cube"},
    )

    result = processor(transition)

    # Check that subtask tokens were added to observation
    observation = result[TransitionKey.OBSERVATION]
    assert OBS_LANGUAGE_SUBTASK_TOKENS in observation
    assert OBS_LANGUAGE_SUBTASK_ATTENTION_MASK in observation

    # Check token structure
    subtask_tokens = observation[OBS_LANGUAGE_SUBTASK_TOKENS]
    subtask_attention_mask = observation[OBS_LANGUAGE_SUBTASK_ATTENTION_MASK]
    assert isinstance(subtask_tokens, torch.Tensor)
    assert isinstance(subtask_attention_mask, torch.Tensor)
    assert subtask_tokens.shape == (8,)
    assert subtask_attention_mask.shape == (8,)
    assert subtask_attention_mask.dtype == torch.bool


@skip_if_package_missing("transformers")
def test_subtask_tokenization_not_added_when_none():
    """Test that subtask tokens are NOT added to observation when subtask is None."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=8)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task"},  # No subtask
    )

    result = processor(transition)

    # Check that subtask tokens were NOT added to observation
    observation = result[TransitionKey.OBSERVATION]
    assert OBS_LANGUAGE_SUBTASK_TOKENS not in observation
    assert OBS_LANGUAGE_SUBTASK_ATTENTION_MASK not in observation

    # But main task tokens should still be present
    assert f"{OBS_LANGUAGE}.tokens" in observation
    assert f"{OBS_LANGUAGE}.attention_mask" in observation


@skip_if_package_missing("transformers")
def test_subtask_tokenization_not_added_when_subtask_value_is_none():
    """Test that subtask tokens are NOT added when subtask value is explicitly None."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=8)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task", "subtask": None},
    )

    result = processor(transition)

    # Check that subtask tokens were NOT added to observation
    observation = result[TransitionKey.OBSERVATION]
    assert OBS_LANGUAGE_SUBTASK_TOKENS not in observation
    assert OBS_LANGUAGE_SUBTASK_ATTENTION_MASK not in observation


@skip_if_package_missing("transformers")
def test_subtask_tokenization_list_of_strings():
    """Test subtask tokenization with list of strings."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=8)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task", "subtask": ["pick up", "place down"]},
    )

    result = processor(transition)

    # Check that subtask tokens were added to observation
    observation = result[TransitionKey.OBSERVATION]
    assert OBS_LANGUAGE_SUBTASK_TOKENS in observation
    assert OBS_LANGUAGE_SUBTASK_ATTENTION_MASK in observation

    # Check token structure for batch
    subtask_tokens = observation[OBS_LANGUAGE_SUBTASK_TOKENS]
    subtask_attention_mask = observation[OBS_LANGUAGE_SUBTASK_ATTENTION_MASK]
    assert subtask_tokens.shape == (2, 8)  # batch_size=2, seq_len=8
    assert subtask_attention_mask.shape == (2, 8)


@skip_if_package_missing("transformers")
def test_subtask_tokenization_device_cpu():
    """Test that subtask tokens are on CPU when other tensors are on CPU."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    # Create transition with CPU tensors
    observation = {OBS_STATE: torch.randn(10)}  # CPU tensor
    action = torch.randn(5)  # CPU tensor
    transition = create_transition(
        observation=observation,
        action=action,
        complementary_data={"task": "main task", "subtask": "pick up cube"},
    )

    result = processor(transition)

    # Check that subtask tokens are on CPU
    subtask_tokens = result[TransitionKey.OBSERVATION][OBS_LANGUAGE_SUBTASK_TOKENS]
    subtask_attention_mask = result[TransitionKey.OBSERVATION][OBS_LANGUAGE_SUBTASK_ATTENTION_MASK]

    assert subtask_tokens.device.type == "cpu"
    assert subtask_attention_mask.device.type == "cpu"


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
@skip_if_package_missing("transformers")
def test_subtask_tokenization_device_cuda():
    """Test that subtask tokens are moved to CUDA when other tensors are on CUDA."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    # Create transition with CUDA tensors
    observation = {OBS_STATE: torch.randn(10).cuda()}  # CUDA tensor
    action = torch.randn(5).cuda()  # CUDA tensor
    transition = create_transition(
        observation=observation,
        action=action,
        complementary_data={"task": "main task", "subtask": "pick up cube"},
    )

    result = processor(transition)

    # Check that subtask tokens are on CUDA
    subtask_tokens = result[TransitionKey.OBSERVATION][OBS_LANGUAGE_SUBTASK_TOKENS]
    subtask_attention_mask = result[TransitionKey.OBSERVATION][OBS_LANGUAGE_SUBTASK_ATTENTION_MASK]

    assert subtask_tokens.device.type == "cuda"
    assert subtask_attention_mask.device.type == "cuda"


@skip_if_package_missing("transformers")
def test_subtask_tokenization_preserves_other_observation_data():
    """Test that subtask tokenization preserves other observation data."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    original_state = torch.tensor([1.0, 2.0, 3.0])
    transition = create_transition(
        observation={"state": original_state.clone()},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task", "subtask": "pick up cube"},
    )

    result = processor(transition)
    observation = result[TransitionKey.OBSERVATION]

    # Check that original observation data is preserved
    assert torch.equal(observation["state"], original_state)

    # Check that both task and subtask tokens are present
    assert f"{OBS_LANGUAGE}.tokens" in observation
    assert f"{OBS_LANGUAGE}.attention_mask" in observation
    assert OBS_LANGUAGE_SUBTASK_TOKENS in observation
    assert OBS_LANGUAGE_SUBTASK_ATTENTION_MASK in observation


@skip_if_package_missing("transformers")
def test_subtask_attention_mask_dtype():
    """Test that subtask attention mask has correct dtype (bool)."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task", "subtask": "pick up cube"},
    )

    result = processor(transition)
    observation = result[TransitionKey.OBSERVATION]

    subtask_attention_mask = observation[OBS_LANGUAGE_SUBTASK_ATTENTION_MASK]
    assert subtask_attention_mask.dtype == torch.bool


@skip_if_package_missing("transformers")
def test_subtask_tokenization_deterministic():
    """Test that subtask tokenization is deterministic for the same input."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=10)

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task", "subtask": "consistent subtask"},
    )

    result1 = processor(transition)
    result2 = processor(transition)

    subtask_tokens1 = result1[TransitionKey.OBSERVATION][OBS_LANGUAGE_SUBTASK_TOKENS]
    subtask_tokens2 = result2[TransitionKey.OBSERVATION][OBS_LANGUAGE_SUBTASK_TOKENS]
    subtask_mask1 = result1[TransitionKey.OBSERVATION][OBS_LANGUAGE_SUBTASK_ATTENTION_MASK]
    subtask_mask2 = result2[TransitionKey.OBSERVATION][OBS_LANGUAGE_SUBTASK_ATTENTION_MASK]

    # Results should be identical
    assert torch.equal(subtask_tokens1, subtask_tokens2)
    assert torch.equal(subtask_mask1, subtask_mask2)


@skip_if_package_missing("transformers")
@patch("lerobot.processor.tokenizer_processor.AutoTokenizer")
def test_subtask_tokenization_integration_with_pipeline(mock_auto_tokenizer):
    """Test subtask tokenization works correctly with DataProcessorPipeline."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    mock_auto_tokenizer.from_pretrained.return_value = mock_tokenizer

    tokenizer_processor = TokenizerProcessorStep(tokenizer_name="test-tokenizer", max_length=6)
    robot_processor = DataProcessorPipeline(
        [tokenizer_processor], to_transition=identity_transition, to_output=identity_transition
    )

    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task", "subtask": "subtask instruction"},
    )

    result = robot_processor(transition)

    # Check that observation exists and both tokenizations were applied
    assert TransitionKey.OBSERVATION in result
    observation = result[TransitionKey.OBSERVATION]

    # Check task tokens
    assert f"{OBS_LANGUAGE}.tokens" in observation
    assert f"{OBS_LANGUAGE}.attention_mask" in observation

    # Check subtask tokens
    assert OBS_LANGUAGE_SUBTASK_TOKENS in observation
    assert OBS_LANGUAGE_SUBTASK_ATTENTION_MASK in observation

    # Check shapes
    assert observation[f"{OBS_LANGUAGE}.tokens"].shape == (6,)
    assert observation[OBS_LANGUAGE_SUBTASK_TOKENS].shape == (6,)


@skip_if_package_missing("transformers")
def test_subtask_not_added_for_unsupported_types():
    """Test that subtask tokens are not added when subtask has unsupported type."""
    mock_tokenizer = MockTokenizer(vocab_size=100)
    processor = TokenizerProcessorStep(tokenizer=mock_tokenizer, max_length=8)

    # Test with integer subtask
    transition = create_transition(
        observation={"state": torch.tensor([1.0, 2.0])},
        action=torch.tensor([0.1, 0.2]),
        complementary_data={"task": "main task", "subtask": 123},
    )

    result = processor(transition)
    observation = result[TransitionKey.OBSERVATION]

    # Subtask tokens should NOT be added for unsupported types
    assert OBS_LANGUAGE_SUBTASK_TOKENS not in observation
    assert OBS_LANGUAGE_SUBTASK_ATTENTION_MASK not in observation

    # But main task tokens should still be present
    assert f"{OBS_LANGUAGE}.tokens" in observation