Revert "use validation instead"

Author: Kovbo

src/art/preprocessing/tokenize.py

@@ -139,14 +139,16 @@ 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):
-            # 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 (
+            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 there are no trainable assistant messages, return None
     if last_assistant_index == -1:
         return None
@@ -187,7 +189,7 @@ def tokenize_trajectory(
                     (
                         message_or_choice
                         if isinstance(message_or_choice, dict)
-                        and message_or_choice["role"] != "assistant"
+                        and not message_or_choice["role"] == "assistant"
                         else {
                             "role": "assistant",
                             "content": sentinal_token,
@@ -203,7 +205,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 message["role"] != "assistant":
+        if isinstance(message, dict) and not message["role"] == "assistant":
             continue
         start = token_ids.index(sentinal_token_id)
         end = start + 1
@@ -212,7 +214,6 @@ 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(
@@ -223,7 +224,6 @@ 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,19 +31,6 @@ 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)
 
@@ -59,19 +46,6 @@ 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()
@@ -123,7 +97,6 @@ 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(
@@ -143,7 +116,7 @@ def get_messages(messages_and_choices: MessagesAndChoices) -> Messages:
                 }
             )
         else:
-            # Regular Message dict
+            # Ensure content is always a string for tokenizer chat templates
             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:
-            inputs.pop("pixel_values", None)
+            del inputs["pixel_values"]  # type: ignore
         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:
-            inputs.pop("image_grid_thw", None)
+            del inputs["image_grid_thw"]  # type: ignore
 
         # Move tensors to the correct device
         inputs = {