use validation instead

Author: Kovbo

src/art/preprocessing/tokenize.py

@@ -139,16 +139,14 @@ def tokenize_trajectory(
     # Find the index of the last assistant message
     last_assistant_index = -1
     for i, message in enumerate(history.messages_and_choices):
-        if (
-            isinstance(message, dict)
-            and message["role"] == "assistant"
-            and allow_training_without_logprobs
-        ):
-            last_assistant_index = i
-        elif not isinstance(message, dict) and (
-            message.logprobs or allow_training_without_logprobs
-        ):
-            last_assistant_index = i
+        if isinstance(message, dict):
+            # Message dict
+            if message["role"] == "assistant" and allow_training_without_logprobs:
+                last_assistant_index = i
+        else:
+            # Choice object
+            if message.logprobs is not None or allow_training_without_logprobs:
+                last_assistant_index = i
     # If there are no trainable assistant messages, return None
     if last_assistant_index == -1:
         return None
@@ -189,7 +187,7 @@ def tokenize_trajectory(
                     (
                         message_or_choice
                         if isinstance(message_or_choice, dict)
-                        and not message_or_choice["role"] == "assistant"
+                        and message_or_choice["role"] != "assistant"
                         else {
                             "role": "assistant",
                             "content": sentinal_token,
@@ -205,7 +203,7 @@ def tokenize_trajectory(
     assistant_mask: list[int] = [0] * len(token_ids)
     logprobs = [float("nan")] * len(token_ids)
     for message in messages_and_choices:
-        if isinstance(message, dict) and not message["role"] == "assistant":
+        if isinstance(message, dict) and message["role"] != "assistant":
             continue
         start = token_ids.index(sentinal_token_id)
         end = start + 1
@@ -214,6 +212,7 @@ def tokenize_trajectory(
         except IndexError:
             end_token_id = None
         if isinstance(message, dict):
+            # Message dict
             content = message.get("content")
             assert isinstance(content, str)
             content_token_ids = tokenizer.encode(
@@ -224,6 +223,7 @@ def tokenize_trajectory(
             logprobs[start:end] = [float("nan")] * len(content_token_ids)
             assistant_mask[start:end] = [1] * len(content_token_ids)
         else:
+            # Choice object
             choice = message
             assert choice.logprobs or allow_training_without_logprobs, (
                 "Chat completion choices must have logprobs"

src/art/trajectories.py

@@ -31,6 +31,19 @@ class History(pydantic.BaseModel):
     messages_and_choices: MessagesAndChoices
     tools: Tools | None = None
 
+    @pydantic.field_validator("messages_and_choices", mode="before")
+    @classmethod
+    def deserialize_choices(cls, v: list[Any]) -> list[Any]:
+        """Convert serialized Choice dicts back to Choice objects."""
+        result = []
+        for item in v:
+            if isinstance(item, dict) and "message" in item and "index" in item:
+                # This is a serialized Choice dict - convert back to Choice object
+                result.append(Choice.model_validate(item))
+            else:
+                result.append(item)
+        return result
+
     def messages(self) -> Messages:
         return get_messages(self.messages_and_choices)
 
@@ -46,6 +59,19 @@ class Trajectory(pydantic.BaseModel):
     logs: list[str] = []
     start_time: datetime = pydantic.Field(default_factory=datetime.now, exclude=True)
 
+    @pydantic.field_validator("messages_and_choices", mode="before")
+    @classmethod
+    def deserialize_choices(cls, v: list[Any]) -> list[Any]:
+        """Convert serialized Choice dicts back to Choice objects."""
+        result = []
+        for item in v:
+            if isinstance(item, dict) and "message" in item and "index" in item:
+                # This is a serialized Choice dict - convert back to Choice object
+                result.append(Choice.model_validate(item))
+            else:
+                result.append(item)
+        return result
+
     def __init__(self, **data: Any):
         super().__init__(**data)
         self.start_time = datetime.now()
@@ -97,6 +123,7 @@ def get_messages(messages_and_choices: MessagesAndChoices) -> Messages:
     messages: Messages = []
     for message_or_choice in messages_and_choices:
         if isinstance(message_or_choice, Choice):
+            # Choice object (always a Choice after Pydantic validation)
             content = message_or_choice.message.content or ""
             tool_calls = message_or_choice.message.tool_calls or []
             messages.append(
@@ -116,7 +143,7 @@ def get_messages(messages_and_choices: MessagesAndChoices) -> Messages:
                 }
             )
         else:
-            # Ensure content is always a string for tokenizer chat templates
+            # Regular Message dict
             msg = dict(message_or_choice)
             if msg.get("content") is None:
                 msg["content"] = ""

src/art/unsloth/train.py

@@ -72,11 +72,11 @@ def compute_loss(
         if inputs.get("pixel_values") and inputs["pixel_values"][0] is not None:
             inputs["pixel_values"] = inputs["pixel_values"][0]  # type: ignore
         else:
-            del inputs["pixel_values"]  # type: ignore
+            inputs.pop("pixel_values", None)
         if inputs.get("image_grid_thw") and inputs["image_grid_thw"][0] is not None:
             inputs["image_grid_thw"] = inputs["image_grid_thw"][0]  # type: ignore
         else:
-            del inputs["image_grid_thw"]  # type: ignore
+            inputs.pop("image_grid_thw", None)
 
         # Move tensors to the correct device
         inputs = {