refactor: Refactor token handling in LocalBackend and TokenizedResult

- Updated LocalBackend to use `token_ids` instead of `tokens` for max token calculation.
- Introduced a cached property in TokenizedResult to dynamically generate `tokens` from `token_ids`.
- Cleaned up assertion formatting in LocalBackend for better readability.

These changes enhance the clarity and efficiency of token management in the codebase.

Author: bradhilton

src/art/local/backend.py

@@ -244,7 +244,7 @@ def _get_packed_tensors(
         )
         if not tokenized_results:
             return None
-        max_tokens = max(len(result.tokens) for result in tokenized_results)
+        max_tokens = max(len(result.token_ids) for result in tokenized_results)
         # Round up max_tokens to the nearest multiple of 2048
         sequence_length = math.ceil(max_tokens / 2048) * 2048
         # Cap sequence length at the model's max sequence length
@@ -416,7 +416,9 @@ def _trajectory_log(self, trajectory: Trajectory) -> str:
             if isinstance(message_or_choice, dict):
                 message = message_or_choice
             else:
-                message = cast(Message, message_or_choice.message.model_dump())  # ty:ignore[possibly-missing-attribute]
+                message = cast(
+                    Message, message_or_choice.message.model_dump()
+                )  # ty:ignore[possibly-missing-attribute]
             formatted_messages.append(format_message(message))
         return header + "\n".join(formatted_messages)
 
@@ -702,9 +704,9 @@ async def _train_model(
             num_gradient_steps = int(
                 result.pop("num_gradient_steps", estimated_gradient_steps)
             )
-            assert num_gradient_steps == estimated_gradient_steps, (
-                f"num_gradient_steps {num_gradient_steps} != estimated_gradient_steps {estimated_gradient_steps}"
-            )
+            assert (
+                num_gradient_steps == estimated_gradient_steps
+            ), f"num_gradient_steps {num_gradient_steps} != estimated_gradient_steps {estimated_gradient_steps}"
             results.append(result)
             yield {**result, "num_gradient_steps": num_gradient_steps}
             pbar.update(1)

src/art/preprocessing/tokenize.py

@@ -1,4 +1,5 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
+from functools import cached_property
 from itertools import takewhile
 import math
 import random
@@ -16,30 +17,34 @@
 class TokenizedResult:
     advantage: float
     chat: str
-    tokens: list[str]
     token_ids: list[int]
     input_pos: list[int]
     assistant_mask: list[int]
     logprobs: list[float]
     pixel_values: torch.Tensor | None
     image_grid_thw: torch.Tensor | None
     trajectory: Trajectory
+    _tokenizer: "PreTrainedTokenizerBase" = field(repr=False, compare=False)
     weight: float = 0.0
     prompt_id: int = 0
     prompt_length: int = 0
 
+    @cached_property
+    def tokens(self) -> list[str]:
+        return [self._tokenizer.decode(token_id) for token_id in self.token_ids]
+
     def without_prompt(self) -> "TokenizedResult":
         return TokenizedResult(
             advantage=self.advantage,
             chat=self.chat,
-            tokens=self.tokens[self.prompt_length :],
             token_ids=self.token_ids[self.prompt_length :],
             input_pos=self.input_pos[self.prompt_length :],
             assistant_mask=self.assistant_mask[self.prompt_length :],
             logprobs=self.logprobs[self.prompt_length :],
             pixel_values=None,
             image_grid_thw=None,
             trajectory=self.trajectory,
+            _tokenizer=self._tokenizer,
             weight=self.weight,
             prompt_id=self.prompt_id,
             prompt_length=0,
@@ -347,14 +352,14 @@ def tokenize_trajectory(
     return TokenizedResult(
         advantage=advantage,
         chat=chat,
-        tokens=[tokenizer.decode(token_id) for token_id in token_ids],
         token_ids=token_ids,
         input_pos=list(range(len(token_ids))),
         assistant_mask=assistant_mask,
         logprobs=logprobs,
         pixel_values=pixel_values,
         image_grid_thw=image_grid_thw,
         trajectory=trajectory,
+        _tokenizer=tokenizer,
     )