[Fix] Keep Thyme sandbox images in user observation turns

tests/test_thyme_sandbox_history.py

@@ -0,0 +1,172 @@
+"""Regression tests for Thyme's iterative sandbox conversation history."""
+
+import copy
+import importlib.util
+import sys
+import types
+
+
+class FakeInputs(dict):
+
+    def __init__(self):
+        super().__init__(input_ids=[[1]])
+        self.input_ids = self["input_ids"]
+
+    def to(self, device):
+        assert device == "cuda"
+        return self
+
+
+class FakeTokenizer:
+    eos_token_id = 99
+
+    def __init__(self):
+        self.outputs = iter(
+            [
+                "<code>make_crop()</code>",
+                "<answer>done</answer>",
+            ]
+        )
+
+    def batch_decode(self, *args, **kwargs):
+        return [next(self.outputs)]
+
+
+class FakeProcessor:
+
+    def __init__(self):
+        self.tokenizer = FakeTokenizer()
+        self.template_calls = []
+
+    def apply_chat_template(self, messages, **kwargs):
+        self.template_calls.append((copy.deepcopy(messages[0]), kwargs))
+        return ["rendered prompt"]
+
+    def __call__(self, **kwargs):
+        return FakeInputs()
+
+
+class FakeModel:
+
+    def __init__(self):
+        self.generate_calls = []
+
+    def generate(self, **kwargs):
+        self.generate_calls.append(kwargs)
+        return [[1, len(self.generate_calls)]]
+
+
+def _load_thyme_module(monkeypatch, sandbox_result, vision_histories):
+    packages = {
+        "vlmeval": "vlmeval",
+        "vlmeval.vlm": "vlmeval/vlm",
+        "vlmeval.vlm.thyme": "vlmeval/vlm/thyme",
+    }
+    for name, path in packages.items():
+        package = types.ModuleType(name)
+        package.__path__ = [path]
+        monkeypatch.setitem(sys.modules, name, package)
+
+    torch = types.ModuleType("torch")
+    monkeypatch.setitem(sys.modules, "torch", torch)
+
+    base = types.ModuleType("vlmeval.vlm.base")
+    base.BaseModel = type("BaseModel", (), {})
+    monkeypatch.setitem(sys.modules, "vlmeval.vlm.base", base)
+
+    prompt = types.ModuleType("vlmeval.vlm.thyme.prompt")
+    prompt.ThymePromptMixin = type("ThymePromptMixin", (), {})
+    monkeypatch.setitem(sys.modules, "vlmeval.vlm.thyme.prompt", prompt)
+
+    sandbox = types.ModuleType("vlmeval.vlm.thyme.sandbox")
+    sandbox.execute_code_in_sandbox = lambda *args, **kwargs: sandbox_result
+    monkeypatch.setitem(sys.modules, "vlmeval.vlm.thyme.sandbox", sandbox)
+
+    utils = types.ModuleType("vlmeval.vlm.thyme.utils")
+    utils.REASONING_SYS_PROMPT = "reason"
+    utils.SIMPLE_SYS_PROMPT = "simple"
+    utils.SPECIAL_STRING_LIST = ["</code>", "</answer>"]
+    utils.generate_prompt_final_qa = lambda question, image: question
+    utils.generate_prompt_simple_qa = lambda question: question
+    monkeypatch.setitem(sys.modules, "vlmeval.vlm.thyme.utils", utils)
+
+    qwen_vl_utils = types.ModuleType("qwen_vl_utils")
+
+    def process_vision_info(messages):
+        vision_histories.append(copy.deepcopy(messages[0]))
+        return [], []
+
+    qwen_vl_utils.process_vision_info = process_vision_info
+    monkeypatch.setitem(sys.modules, "qwen_vl_utils", qwen_vl_utils)
+
+    spec = importlib.util.spec_from_file_location(
+        "vlmeval.vlm.thyme.model", "vlmeval/vlm/thyme/model.py"
+    )
+    module = importlib.util.module_from_spec(spec)
+    monkeypatch.setitem(sys.modules, "vlmeval.vlm.thyme.model", module)
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_sandbox_image_is_a_user_observation_in_reencoded_history(
+        monkeypatch, tmp_path):
+    original = tmp_path / "original.png"
+    crop = tmp_path / "crop.png"
+    original.touch()
+    crop.touch()
+    vision_histories = []
+    sandbox_result = ([str(crop)], "", "", {})
+    module = _load_thyme_module(monkeypatch, sandbox_result, vision_histories)
+
+    thyme = module.Thyme.__new__(module.Thyme)
+    thyme.min_pixels = None
+    thyme.max_pixels = None
+    thyme.fps = 2.0
+    thyme.nframe = 64
+    thyme.FRAME_FACTOR = 2
+    thyme.verbose = False
+    thyme.max_retry = 1
+    thyme.max_iterations = 2
+    thyme.post_process = True
+    thyme.temperature = 0.01
+    thyme.generate_kwargs = {
+        "temperature": thyme.temperature,
+        "stop_strings": ["</code>", "</answer>"],
+    }
+    thyme.processor = FakeProcessor()
+    thyme.model = FakeModel()
+
+    answer = thyme.generate_inner_transformers(
+        [
+            {"type": "image", "value": str(original)},
+            {"type": "text", "value": "What is shown?"},
+        ],
+        temp_output_dir=str(tmp_path),
+    )
+
+    assert answer == "done"
+    assert len(thyme.processor.template_calls) == 2
+    assert all(
+        kwargs["add_generation_prompt"]
+        for _, kwargs in thyme.processor.template_calls
+    )
+    assert all("past_key_values" not in call for call in thyme.model.generate_calls)
+
+    second_history = thyme.processor.template_calls[1][0]
+    assert [message["role"] for message in second_history] == [
+        "system",
+        "user",
+        "assistant",
+        "user",
+    ]
+    assistant_items = second_history[2]["content"]
+    assert assistant_items == [
+        {"type": "text", "text": "<code>make_crop()</code>"}
+    ]
+    observation_items = second_history[3]["content"]
+    assert observation_items == [
+        {"type": "text", "text": "<sandbox_output>"},
+        {"type": "image", "image": str(crop)},
+        {"type": "text", "text": "</sandbox_output>"},
+    ]
+    assert vision_histories == [call[0] for call in thyme.processor.template_calls]

vlmeval/vlm/thyme/model.py

@@ -6,7 +6,6 @@
 import warnings
 
 import torch
-from transformers.cache_utils import DynamicCache
 
 from ..base import BaseModel
 from .prompt import ThymePromptMixin
@@ -296,9 +295,6 @@ def generate_inner_transformers(
             # maybe perform code execution.
             conversation_history = copy.deepcopy(messages)
 
-            # For each generation, we initialize a KV-Cache to speed up
-            # inference.
-            kv_cache = DynamicCache()
             # Maintain a dictionary to save context (local & global vars.) for
             # code execution.
             previous_execution_context = {}
@@ -313,19 +309,15 @@ def generate_inner_transformers(
             # execution.
             while retry_iterations > 0:
                 retry_iterations -= 1
-                generated_content = []
+                assistant_content = []
+                sandbox_content = []
                 if self.verbose:
                     print(
                         f"\033[32m\n--- Iteration {self.max_iterations - retry_iterations} ---\033[0m"
                     )
 
                 text = self.processor.apply_chat_template(
-                    [conversation_history], tokenize=False, add_generation_prompt=(
-                        retry_iterations == self.max_iterations - 1), )
-
-                if retry_iterations != self.max_iterations - 1:
-                    if text[0].endswith("<|im_end|>\n"):
-                        text[0] = text[0][: -len("<|im_end|>\n")]
+                    [conversation_history], tokenize=False, add_generation_prompt=True)
                 images, videos = process_vision_info([conversation_history])
                 inputs = self.processor(
                     text=text,
@@ -336,16 +328,13 @@ def generate_inner_transformers(
                 )
                 inputs = inputs.to("cuda")
 
-                # just in case this iteration is invalid, we need to roll back,
-                # thus making a backup.
-                last_kv_cache = copy.deepcopy(kv_cache)
                 # bkup context. roll back when we fail to execute the generated
                 # code.
                 last_execution_context = copy.deepcopy(
                     self._remove_unpickable_values(previous_execution_context)
                 )
                 generated_ids = self.model.generate(
-                    **inputs, **self.generate_kwargs, past_key_values=kv_cache
+                    **inputs, **self.generate_kwargs
                 )
                 generated_ids = [
                     output_ids[len(input_ids):]
@@ -360,7 +349,7 @@ def generate_inner_transformers(
 
                 # Case 1: directly give answer
                 if "</answer>" in generated_text_segment:
-                    generated_content.append(
+                    assistant_content.append(
                         {"type": "text", "text": generated_text_segment},
                     )
 
@@ -398,14 +387,16 @@ def generate_inner_transformers(
                     previous_execution_context = current_execution_context
                     if not processed_img_paths:
                         # deemed as unsuccessful iteration. roll back status.
-                        kv_cache = last_kv_cache
                         previous_execution_context = last_execution_context
                         print(f"{error_msg}")
                         continue
 
                     has_valid_images = False
-                    generated_content += [
-                        {"type": "text", "text": generated_text_segment},
+                    if not assistant_content:
+                        assistant_content.append(
+                            {"type": "text", "text": generated_text_segment}
+                        )
+                    sandbox_content += [
                         {"type": "text", "text": "<sandbox_output>"},
                     ]
                     first_path = processed_img_paths[0]
@@ -417,15 +408,15 @@ def generate_inner_transformers(
                                 # output block
                                 if not has_valid_images:
                                     has_valid_images = True
-                                generated_content.append(
+                                sandbox_content.append(
                                     {"type": "image", "image": img_path}
                                 )
                     else:
-                        generated_content.append(
+                        sandbox_content.append(
                             {"type": "text", "text": first_path})
 
                     if has_valid_images or not os.path.exists(first_path):
-                        generated_content.append(
+                        sandbox_content.append(
                             {"type": "text", "text": "</sandbox_output>"}
                         )
                     else:
@@ -445,17 +436,19 @@ def generate_inner_transformers(
                         self.generate_kwargs["temperature"] = 1.0
                         break
 
-                # Update conversation_history with the latest generated segment
-                # If the last message was 'user', start a new 'assistant'
-                # message
-                if conversation_history[-1]["role"] == "user":
+                # Keep model output in the assistant turn and sandbox output in
+                # a separate user observation turn. The next full-history
+                # encoding can then align every image with a user placeholder.
+                if assistant_content and conversation_history[-1]["role"] == "user":
+                    conversation_history.append(
+                        {"role": "assistant", "content": assistant_content}
+                    )
+                elif assistant_content and conversation_history[-1]["role"] == "assistant":
+                    conversation_history[-1]["content"] += assistant_content
+                if sandbox_content:
                     conversation_history.append(
-                        {"role": "assistant", "content": generated_content}
+                        {"role": "user", "content": sandbox_content}
                     )
-                # If the last message was 'assistant', append to its last text
-                # content item
-                elif conversation_history[-1]["role"] == "assistant":
-                    conversation_history[-1]["content"] += generated_content
 
                 # --- Check for final answer tag if no code was processed in this segment ---
                 if "</answer>" in generated_text_segment: