chore: Refactor message handling and enhance tokenization logic

- Updated imports for clarity and consistency in `types.py`.
- Added `drop_zero_advantage_trajectories` parameter to `tokenize_trajectory_groups` function in `tokenize.py` to control trajectory filtering.
- Introduced `_normalize_message_or_choice` function in `server.py` to standardize message validation and conversion.
- Enhanced `_message_or_choice_to_dict` function in `client.py` to utilize the new message adapter for improved validation.

Author: bradhilton

src/art/preprocessing/tokenize.py

@@ -155,6 +155,7 @@ def tokenize_trajectory_groups(
     allow_training_without_logprobs: bool,
     scale_rewards: bool,
     shuffle_group_trajectories: bool = True,
+    drop_zero_advantage_trajectories: bool = True,
     image_processor: BaseImageProcessor | None = None,
 ) -> Generator["TokenizedResult", None, None]:
     for group in trajectory_groups:
@@ -172,8 +173,7 @@ def tokenize_trajectory_groups(
             advantage = trajectory.reward - reward_mean
             if scale_rewards:
                 advantage /= reward_std + 1e-6
-            # Skip trajectories with no advantage
-            if advantage == 0:
+            if advantage == 0 and drop_zero_advantage_trajectories:
                 continue
             trajectory_results: list[TokenizedResult] = []
             for history in [

src/art/tinker/client.py

@@ -12,17 +12,23 @@
 from openai.resources.models import AsyncModels
 from openai.types import Model
 from openai.types.chat.chat_completion import Choice
+from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam
 from openai.types.completion_usage import CompletionUsage
+from pydantic import TypeAdapter
 
 from art.types import Message, MessageOrChoice, MessagesAndChoices, Tools
 
 ParsedMessageOrChoice = Choice | Message
 ParsedMessagesAndChoices = list[ParsedMessageOrChoice]
+_MESSAGE_ADAPTER = TypeAdapter(ChatCompletionMessageParam)
 
 
 def _message_or_choice_to_dict(message_or_choice: MessageOrChoice) -> dict[str, Any]:
     if isinstance(message_or_choice, dict):
-        return cast(dict[str, Any], message_or_choice)
+        validated = _MESSAGE_ADAPTER.validate_python(message_or_choice)
+        return cast(
+            dict[str, Any], _MESSAGE_ADAPTER.dump_python(validated, mode="json")
+        )
     if isinstance(message_or_choice, BaseModel):
         return cast(dict[str, Any], message_or_choice.to_dict())
     to_dict = getattr(message_or_choice, "to_dict", None)

src/art/tinker/server.py

@@ -26,7 +26,7 @@
 )
 from openai.types.chat.completion_create_params import CompletionCreateParams
 from openai.types.completion_usage import CompletionUsage
-from pydantic import BaseModel, Field, SkipValidation
+from pydantic import BaseModel, Field, SkipValidation, TypeAdapter
 import tinker
 from transformers.tokenization_utils_base import BatchEncoding
 import uvicorn
@@ -49,6 +49,7 @@ class ModelUpsert(BaseModel):
 
 
 WireMessagesAndChoices = list[Choice | Message]
+_MESSAGE_ADAPTER = TypeAdapter(ChatCompletionMessageParam)
 
 
 class MessagesAndChoicesWithLogprobsArgs(BaseModel):
@@ -63,6 +64,14 @@ class MessagesAndChoicesWithLogprobs(BaseModel):
     usages: list[CompletionUsage]
 
 
+def _normalize_message_or_choice(
+    message_or_choice: Choice | Message,
+) -> Choice | Message:
+    if isinstance(message_or_choice, Choice):
+        return message_or_choice
+    return cast(Message, _MESSAGE_ADAPTER.validate_python(message_or_choice))
+
+
 def _normalize_qwen3_5_messages(
     base_model: str, messages: list[ChatCompletionMessageParam]
 ) -> list[dict[str, Any]]:
@@ -264,7 +273,10 @@ async def add_logprobs(model: str, alias: str | None) -> CompletionUsage:
                 ]
             )
             return MessagesAndChoicesWithLogprobs(
-                messages_and_choices=args.messages_and_choices,
+                messages_and_choices=[
+                    _normalize_message_or_choice(item)
+                    for item in args.messages_and_choices
+                ],
                 usages=usages,
             )
 

src/art/types.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass, field
 from typing import Annotated, Literal
 
-from openai.types.chat.chat_completion import Choice
+from openai.types.chat.chat_completion import Choice as Choice
 from openai.types.chat.chat_completion_message_param import ChatCompletionMessageParam
 from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam
 import pydantic