action_token_id_mapping

action_token_id_mapping

Mappings between action-local token IDs and model token IDs.

ActionTokenIdMapping

Bases: ABC

Maps discretizer-local action IDs to the model token-id space.

model_token_count abstractmethod

model_token_count(action_token_count)

Return the model token count before the action tokenizer adds EOS.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

@abstractmethod
def model_token_count(self, action_token_count: int) -> int:
    """Return the model token count before the action tokenizer adds EOS."""

tokenizer_vocab_size

tokenizer_vocab_size(action_token_count)

Return the action tokenizer vocabulary size including EOS.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def tokenizer_vocab_size(self, action_token_count: int) -> int:
    """Return the action tokenizer vocabulary size including EOS."""
    return self.model_token_count(action_token_count) + 1

eos_token_id

eos_token_id(action_token_count)

Return the EOS token ID used after encoded action tokens.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def eos_token_id(self, action_token_count: int) -> int:
    """Return the EOS token ID used after encoded action tokens."""
    return self.model_token_count(action_token_count)

encode abstractmethod

encode(local_token_ids)

Map local action IDs to model token IDs.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

@abstractmethod
def encode(self, local_token_ids: list[int] | np.ndarray) -> np.ndarray:
    """Map local action IDs to model token IDs."""

decode abstractmethod

decode(model_token_ids)

Map model token IDs back to local action IDs.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

@abstractmethod
def decode(self, model_token_ids: np.ndarray | torch.Tensor) -> np.ndarray:
    """Map model token IDs back to local action IDs."""

state_dict abstractmethod

state_dict()

Return serializable state.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

@abstractmethod
def state_dict(self) -> dict[str, Any]:
    """Return serializable state."""

load_state_dict abstractmethod

load_state_dict(state_dict)

Load serializable state.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

@abstractmethod
def load_state_dict(self, state_dict: dict[str, Any]) -> None:
    """Load serializable state."""

save_pretrained

save_pretrained(path)

Save optional external assets.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def save_pretrained(self, path: Path) -> None:
    """Save optional external assets."""
    del path

load_pretrained_assets

load_pretrained_assets(path)

Load optional external assets.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def load_pretrained_assets(self, path: Path) -> None:
    """Load optional external assets."""
    del path

IdentityActionTokenIdMapping

Bases: ActionTokenIdMapping

Use action IDs directly as model token IDs.

model_token_count

model_token_count(action_token_count)

Return the unchanged action-token count.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def model_token_count(self, action_token_count: int) -> int:
    """Return the unchanged action-token count."""
    return action_token_count

encode

encode(local_token_ids)

Return local action IDs unchanged as model token IDs.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def encode(self, local_token_ids: list[int] | np.ndarray) -> np.ndarray:
    """Return local action IDs unchanged as model token IDs."""
    return np.asarray(local_token_ids)

decode

decode(model_token_ids)

Return model token IDs unchanged as local action IDs.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def decode(self, model_token_ids: np.ndarray | torch.Tensor) -> np.ndarray:
    """Return model token IDs unchanged as local action IDs."""
    if isinstance(model_token_ids, torch.Tensor):
        return model_token_ids.detach().cpu().numpy()
    return np.asarray(model_token_ids)

state_dict

state_dict()

Return serializable identity mapping state.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def state_dict(self) -> dict[str, Any]:
    """Return serializable identity mapping state."""
    return {"type": ActionTokenIdMappingType.IDENTITY.value}

load_state_dict

load_state_dict(state_dict)

Load identity mapping state.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def load_state_dict(self, state_dict: dict[str, Any]) -> None:
    """Load identity mapping state."""
    del state_dict

LanguageVocabularyActionTokenIdMapping

LanguageVocabularyActionTokenIdMapping(language_tokenizer_model, num_special_tokens_to_skip=128)

Bases: ActionTokenIdMapping

Place action IDs in the tail of a language tokenizer's token-id space.

Initialize mapping into the tail of a language tokenizer vocabulary.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def __init__(
    self,
    language_tokenizer_model: str,
    num_special_tokens_to_skip: int = 128,
):
    """Initialize mapping into the tail of a language tokenizer vocabulary."""
    self.language_tokenizer_model = language_tokenizer_model
    self.num_special_tokens_to_skip = num_special_tokens_to_skip
    self.language_tokenizer = load_huggingface_tokenizer(
        tokenizer_model=language_tokenizer_model
    )
    if self.language_tokenizer.pad_token is None:
        self.language_tokenizer.pad_token = self.language_tokenizer.eos_token

model_token_count

model_token_count(action_token_count)

Return the language tokenizer vocabulary size if action IDs fit.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def model_token_count(self, action_token_count: int) -> int:
    """Return the language tokenizer vocabulary size if action IDs fit."""
    required_token_count = action_token_count + self.num_special_tokens_to_skip
    if self.language_tokenizer.vocab_size < required_token_count:
        raise ValueError(
            "Language tokenizer token count "
            f"({self.language_tokenizer.vocab_size}) is too small to hold "
            f"action tokens ({action_token_count}) plus skipped special tokens "
            f"({self.num_special_tokens_to_skip}). Required: {required_token_count}"
        )
    self._validate_eos_does_not_overlap_action_tokens(
        action_token_count=action_token_count
    )
    return self.language_tokenizer.vocab_size

tokenizer_vocab_size

tokenizer_vocab_size(action_token_count)

Return the language tokenizer vocabulary size without adding EOS.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def tokenizer_vocab_size(self, action_token_count: int) -> int:
    """Return the language tokenizer vocabulary size without adding EOS."""
    return self.model_token_count(action_token_count=action_token_count)

eos_token_id

eos_token_id(action_token_count)

Return the native language-tokenizer EOS ID.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def eos_token_id(self, action_token_count: int) -> int:
    """Return the native language-tokenizer EOS ID."""
    self.model_token_count(action_token_count=action_token_count)
    eos_token_id = self.language_tokenizer.eos_token_id
    if eos_token_id is None:
        raise ValueError(
            "Language tokenizer must define eos_token_id when used for "
            "action-token EOS."
        )
    return int(eos_token_id)

encode

encode(local_token_ids)

Map local action IDs to language-token tail IDs.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def encode(self, local_token_ids: list[int] | np.ndarray) -> np.ndarray:
    """Map local action IDs to language-token tail IDs."""
    local_tokens = np.asarray(local_token_ids)
    return (
        self.language_tokenizer.vocab_size
        - 1
        - self.num_special_tokens_to_skip
        - local_tokens
    )

decode

decode(model_token_ids)

Map language-token tail IDs back to local action IDs.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def decode(self, model_token_ids: np.ndarray | torch.Tensor) -> np.ndarray:
    """Map language-token tail IDs back to local action IDs."""
    if isinstance(model_token_ids, torch.Tensor):
        model_token_ids = model_token_ids.detach().cpu().numpy()
    return (
        self.language_tokenizer.vocab_size
        - 1
        - self.num_special_tokens_to_skip
        - np.asarray(model_token_ids)
    )

state_dict

state_dict()

Return serializable language-vocabulary mapping state.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def state_dict(self) -> dict[str, Any]:
    """Return serializable language-vocabulary mapping state."""
    return {
        "type": ActionTokenIdMappingType.LANGUAGE_VOCABULARY.value,
        "language_tokenizer_model": self.language_tokenizer_model,
        "num_special_tokens_to_skip": self.num_special_tokens_to_skip,
    }

load_state_dict

load_state_dict(state_dict)

Load language-vocabulary mapping state.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def load_state_dict(self, state_dict: dict[str, Any]) -> None:
    """Load language-vocabulary mapping state."""
    tokenizer_model = state_dict["language_tokenizer_model"]
    needs_reload = tokenizer_model != self.language_tokenizer_model
    self.language_tokenizer_model = tokenizer_model
    self.num_special_tokens_to_skip = state_dict["num_special_tokens_to_skip"]
    if needs_reload:
        self.language_tokenizer = load_huggingface_tokenizer(
            tokenizer_model=tokenizer_model
        )
        if self.language_tokenizer.pad_token is None:
            self.language_tokenizer.pad_token = self.language_tokenizer.eos_token

save_pretrained

save_pretrained(path)

Save the language tokenizer assets.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def save_pretrained(self, path: Path) -> None:
    """Save the language tokenizer assets."""
    self.language_tokenizer.save_pretrained(path / "language_tokenizer")

load_pretrained_assets

load_pretrained_assets(path)

Load saved language tokenizer assets when present.

Source code in src/versatil/data/tokenization/action_token_id_mapping.py

def load_pretrained_assets(self, path: Path) -> None:
    """Load saved language tokenizer assets when present."""
    language_path = path / "language_tokenizer"
    if language_path.exists():
        self.language_tokenizer = load_huggingface_tokenizer(
            tokenizer_model=language_path
        )