Track trajectory finish state using metadata

src/art/preprocessing/tokenize.py

@@ -158,6 +158,12 @@ def tokenize_trajectory(
     Tokenizes a trajectory and returns a TokenizedResult.
     """
     # Find the index of the last assistant message
+
+    # Check if the trajectory is finished
+    continue_final_message = (
+        False if trajectory.metadata.get("finished", False) is True else True
+    )
+
     last_assistant_index = -1
     for i, message in enumerate(history.messages_and_choices):
         if (
@@ -185,7 +191,7 @@ def tokenize_trajectory(
         tokenizer.apply_chat_template(
             cast(list[dict], messages),
             tools=tools,
-            continue_final_message=True,
+            continue_final_message=continue_final_message,
             tokenize=False,
         ),
     )
@@ -194,7 +200,7 @@ def tokenize_trajectory(
         tokenizer.apply_chat_template(
             cast(list[dict], messages),
             tools=tools,
-            continue_final_message=True,
+            continue_final_message=continue_final_message,
         ),
     )
     sentinal_token_id = max(
@@ -229,7 +235,7 @@ def tokenize_trajectory(
         tokenizer.apply_chat_template(
             cast(list[dict], token_template_messages),
             tools=tools,
-            continue_final_message=True,
+            continue_final_message=continue_final_message,
         ),
     )
     assistant_mask: list[int] = [0] * len(token_ids)

src/art/trajectories.py

@@ -59,6 +59,7 @@ def log(self, message: str) -> None:
     def finish(self) -> "Trajectory":
         duration = (datetime.now() - self.start_time).total_seconds()
         self.metrics["duration"] = duration
+        self.metadata["finished"] = True
         return self
 
     @asynccontextmanager