feat: Expand importance_sampling_level options and improve type hints in tokenization

Author: bradhilton

src/art/dev/train.py

@@ -13,7 +13,9 @@ class TrainConfig(TypedDict, total=False):
     epsilon_high: (
         float | None
     )  # asymmetric clip upper bound. Defaults to epsilon when None
-    importance_sampling_level: Literal["token", "sequence"]
+    importance_sampling_level: Literal[
+        "token", "sequence", "average", "harmonic_average"
+    ]
     logprob_calculation_chunk_size: int
     max_negative_advantage_importance_sampling_weight: float
     num_trajectories_learning_rate_multiplier_power: float

src/art/preprocessing/tokenize.py

@@ -163,7 +163,7 @@ def tokenize_trajectory(
         str,
         tokenizer.apply_chat_template(
             cast(list[dict], messages),
-            tools=tools,
+            tools=tools,  # type: ignore
             continue_final_message=True,
             tokenize=False,
         ),
@@ -172,7 +172,7 @@ def tokenize_trajectory(
         list[int],
         tokenizer.apply_chat_template(
             cast(list[dict], messages),
-            tools=tools,
+            tools=tools,  # type: ignore
             continue_final_message=True,
         ),
     )
@@ -198,7 +198,7 @@ def tokenize_trajectory(
                     for message_or_choice in messages_and_choices
                 ],
             ),
-            tools=tools,
+            tools=tools,  # type: ignore
             continue_final_message=True,
         ),
     )

src/art/serverless/backend.py

@@ -133,7 +133,9 @@ async def _train_model(
                 epsilon_high=dev_config.get("epsilon_high"),
                 importance_sampling_level=dev_config.get("importance_sampling_level"),
                 learning_rate=config.learning_rate,
-                max_negative_advantage_importance_sampling_weight=dev_config.get("max_negative_advantage_importance_sampling_weight"),
+                max_negative_advantage_importance_sampling_weight=dev_config.get(
+                    "max_negative_advantage_importance_sampling_weight"
+                ),
                 ppo=dev_config.get("ppo"),
                 precalculate_logprobs=dev_config.get("precalculate_logprobs"),
                 scale_rewards=dev_config.get("scale_rewards"),

src/art/serverless/client.py

@@ -53,7 +53,9 @@ class ExperimentalTrainingConfig(TypedDict, total=False):
     advantage_balance: float | None
     epsilon: float | None
     epsilon_high: float | None
-    importance_sampling_level: Literal["token", "sequence"] | None
+    importance_sampling_level: (
+        Literal["token", "sequence", "average", "harmonic_average"] | None
+    )
     learning_rate: float | None
     max_negative_advantage_importance_sampling_weight: float | None
     ppo: bool | None