tensorrt_llm/inputs/utils.py (2)

393-405: Add unit tests for llava_next placeholder reconstruction (0/1/N images).

Recommend tests that assert:

• Content shape is [{"type":"text",...}, {"type":"image"}, ...] with correct image count.
• No KeyError when there are zero images.
• Works when placeholder is <image> or <|image|> (parametrize via registry stub/mocking).

I can draft pytest cases targeting handle_placeholder_exceptions and the default_multimodal_input_loader → apply_chat_template path. Want me to open a test PR?

---

396-404: Use the registry-provided placeholder and safe .get access

To avoid hard-coding “<image>” (and risking a KeyError if it ever changes or is missing), pull the placeholder string from the multimodal registry and fall back to 0 when absent. Also replace the backslash continuation with parentheses:

-    if model_type == "llava_next":
-        # we need to convert the flattened content back to conversation format
-        for conv, mm_placeholder_count in zip(conversation,
-                                              mm_placeholder_counts):
-            conv["content"] = [{"type": "text", "text": conv["content"]}, \
-                *[{"type": "image"} for _ in range(mm_placeholder_count['<image>'])]]
+    if model_type == "llava_next":
+        # convert the flattened content back to conversation format
+        for conv, mm_placeholder_count in zip(conversation, mm_placeholder_counts):
+            # lookup the exact image placeholder (e.g. "<image>") and default to 0
+            placeholder = MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder(model_type, "image")
+            image_count = mm_placeholder_count.get(placeholder, 0)
+            conv["content"] = (
+                [{"type": "text", "text": conv["content"]}]
+                + [{"type": "image"} for _ in range(image_count)]
+            )

This ensures the code stays in sync with whatever placeholder the registry defines for llava_next and never KeyErrors if counts are empty.

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tensorrt_llm/inputs/utils.py (2)

433-437: Correct arg shape fix for placeholder exceptions (llava_next) — good catch.

Wrapping mm_placeholder_counts in a list aligns with handle_placeholder_exceptions(..., mm_placeholder_counts: list[dict[str, int]]) and prevents zipping a dict (iterating keys) with conversation. This unblocks correct per-message handling for the llava_next path.

---

433-437: It looks like we’re unable to locate and inspect parse_chat_messages_coroutines in tensorrt_llm/serve/chat_utils.py. Without confirming how mm_placeholder_counts is structured (dict vs list) after parsing, we can’t be certain our padding/length-guard logic will cover all scenarios or might even break existing behavior. Please verify manually that:

• parse_chat_messages_coroutines indeed returns mm_placeholder_counts as a single dict (not a list) for PLACEHOLDER_EXCEPTIONS models.
• Multi-message conversations are never passed to apply_chat_template when model_type is in PLACEHOLDER_EXCEPTIONS (currently only "llava_next").
• If lists are possible, adapt the guard accordingly to avoid wrapping an already-list in another list.

Once confirmed, you can safely apply the suggested non-breaking padding guard.