diff --git a/fern/versions/main/pages/curate-video/process-data/captions-preview.mdx b/fern/versions/main/pages/curate-video/process-data/captions-preview.mdx index 7fca23d56a..f1c3b082c6 100644 --- a/fern/versions/main/pages/curate-video/process-data/captions-preview.mdx +++ b/fern/versions/main/pages/curate-video/process-data/captions-preview.mdx @@ -20,6 +20,7 @@ Use the same `model_variant` in `CaptionPreparationStage` and `CaptionGeneration | --- | --- | --- | | `qwen2.5` | `Qwen/Qwen2.5-VL-7B-Instruct` | Default; BF16 unless `fp8=True` or `--captioning-use-fp8-weights` is set | | `qwen3` | `Qwen/Qwen3-VL-8B-Instruct` | Qwen3-VL; BF16 unless FP8 is enabled | +| `qwen3.5` | `Qwen/Qwen3.5-9B` | Qwen3.5; BF16 unless FP8 is enabled | | `nemotron`, `nemotron-bf16` | `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16` | `nemotron` is an alias for the BF16 checkpoint | | `nemotron-fp8` | `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-FP8` | Pre-quantized FP8 checkpoint | | `nemotron-nvfp4` | `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-NVFP4-QAD` | NVFP4 quantization-aware-distilled checkpoint | @@ -84,7 +85,7 @@ python tutorials/video/getting-started/video_split_clip_example.py \ --preview-target-height 240 ``` -`--captioning-algorithm` accepts `qwen2.5`, `qwen3`, `nemotron`, `nemotron-bf16`, `nemotron-fp8`, `nemotron-nvfp4`, and `nemotron-3-nano-omni`. +`--captioning-algorithm` accepts `qwen2.5`, `qwen3`, `qwen3.5`, `nemotron`, `nemotron-bf16`, `nemotron-fp8`, `nemotron-nvfp4`, and `nemotron-3-nano-omni`. @@ -109,7 +110,7 @@ prep = CaptionPreparationStage( | Parameter | Type | Default | Description | | --- | --- | --- | --- | -| `model_variant` | str | `"qwen2.5"` | One of the seven variants in the model table. Must match the generation stage. | +| `model_variant` | str | `"qwen2.5"` | One of the eight variants in the model table. Must match the generation stage. | | `prompt_variant` | `"default"`, `"av"`, `"av-surveillance"` | `"default"` | Built-in caption prompt used when `prompt_text` is not set. | | `prompt_text` | str \| None | `None` | Custom prompt that overrides `prompt_variant`. | | `sampling_fps` | float | 2.0 | Frames per second sampled from the source clip. | @@ -140,7 +141,7 @@ gen = CaptionGenerationStage( | Parameter | Type | Default | Description | | --- | --- | --- | --- | | `model_dir` | str | `"models/qwen"` | Base directory under which checkpoint-specific directories are created. | -| `model_variant` | str | `"qwen2.5"` | One of the seven variants in the model table. | +| `model_variant` | str | `"qwen2.5"` | One of the eight variants in the model table. | | `caption_batch_size` | int | 16 | Generation batch size. The example CLI defaults to 8. | | `fp8` | bool | `False` | Quantize Qwen weights to FP8. Select `nemotron-fp8` for the Nemotron FP8 checkpoint. | | `max_output_tokens` | int | 512 | Maximum tokens generated for each caption. | @@ -171,9 +172,9 @@ enhance = CaptionEnhancementStage( ) ``` -The enhancement model can be `qwen2.5` (`Qwen/Qwen2.5-14B-Instruct`) or `qwen3` (`Qwen/Qwen3-14B`). `captioning_model_variant` must identify the caption key produced by the earlier generation stage; it does not have to match the enhancement model. +The enhancement model can be `qwen2.5` (`Qwen/Qwen2.5-14B-Instruct`), `qwen3` (`Qwen/Qwen3-14B`), or `qwen3.5` (`Qwen/Qwen3.5-27B`). Qwen3.5 has no 14B dense checkpoint, so `qwen3.5` uses the 27B one and runs with thinking disabled so enhanced captions exclude reasoning traces. At roughly 54 GB in BF16 it is much larger than the other two enhancement models: enable `--enhance-captions-use-fp8-weights` to fit it on a GPU with less than 80 GB. Note this differs from the `qwen3.5` captioning checkpoint (`Qwen/Qwen3.5-9B`), so selecting `qwen3.5` for both stages downloads two models. `captioning_model_variant` must identify the caption key produced by the earlier generation stage; it does not have to match the enhancement model. -In the example script, `--enhance-captions-algorithm` selects `qwen2.5` or `qwen3`. The separate `--enhanced-caption-models` output selector currently accepts only `qwen_lm`. +In the example script, `--enhance-captions-algorithm` selects `qwen2.5`, `qwen3`, or `qwen3.5`. The separate `--enhanced-caption-models` output selector currently accepts only `qwen_lm`. ```bash python tutorials/video/getting-started/video_split_clip_example.py \ @@ -188,7 +189,7 @@ python tutorials/video/getting-started/video_split_clip_example.py \ | Parameter | Type | Default | Description | | --- | --- | --- | --- | | `model_dir` | str | `"models/qwen"` | Base directory for Qwen language-model weights. | -| `model_variant` | `"qwen2.5"`, `"qwen3"` | `"qwen2.5"` | Text-only model used for enhancement. | +| `model_variant` | `"qwen2.5"`, `"qwen3"`, `"qwen3.5"` | `"qwen2.5"` | Text-only model used for enhancement. | | `captioning_model_variant` | str | `"qwen2.5"` | Key in `window.caption` to read. Set it to the earlier captioning variant. | | `prompt_variant` | `"default"`, `"av-surveillance"` | `"default"` | Built-in enhancement prompt used when `prompt_text` is not set. | | `prompt_text` | str \| None | `None` | Custom enhancement system prompt. | diff --git a/nemo_curator/models/prompt_formatter.py b/nemo_curator/models/prompt_formatter.py index 737959e16e..6f894d8576 100644 --- a/nemo_curator/models/prompt_formatter.py +++ b/nemo_curator/models/prompt_formatter.py @@ -29,6 +29,7 @@ VARIANT_MAPPING: dict[str, str] = { "qwen2.5": "Qwen/Qwen2.5-VL-7B-Instruct", "qwen3": "Qwen/Qwen3-VL-8B-Instruct", + "qwen3.5": "Qwen/Qwen3.5-9B", **_NEMOTRON_VARIANTS_INFO, "nemotron-3-nano-omni": _NEMOTRON_3_NANO_OMNI_HF_ID, } @@ -80,7 +81,7 @@ def generate_inputs( - "multi_modal_data": Dictionary containing processed "video" inputs """ - if self.prompt_variant in {"qwen2.5", "qwen3"}: + if self.prompt_variant in {"qwen2.5", "qwen3", "qwen3.5"}: return self._generate_qwen_inputs(prompt, video_inputs, override_text_prompt, fps) if self.prompt_variant.startswith("nemotron"): @@ -99,10 +100,14 @@ def _generate_qwen_inputs( """Generate inputs for Qwen models.""" message = self._create_qwen_message(prompt) if self.text_prompt is None or override_text_prompt: + template_kwargs: dict[str, Any] = {} + if self.prompt_variant == "qwen3.5": + template_kwargs["enable_thinking"] = False self.text_prompt = self.processor.apply_chat_template( message, tokenize=False, add_generation_prompt=True, + **template_kwargs, ) video_data = video_inputs if video_inputs is not None: diff --git a/nemo_curator/models/qwen_lm.py b/nemo_curator/models/qwen_lm.py index 3424a45a15..06f81f406a 100644 --- a/nemo_curator/models/qwen_lm.py +++ b/nemo_curator/models/qwen_lm.py @@ -38,9 +38,12 @@ class SamplingParams: from nemo_curator.models.base import ModelInterface +_QWEN_LM_TEXT_ONLY_MULTIMODAL_VARIANTS: Final = frozenset({"qwen3.5"}) + _QWEN_LM_VARIANTS_INFO: Final = { "qwen2.5": ("Qwen/Qwen2.5-14B-Instruct", "cf98f3b"), "qwen3": ("Qwen/Qwen3-14B", "f8c293d"), + "qwen3.5": ("Qwen/Qwen3.5-27B", "fc05dae"), } @@ -90,10 +93,13 @@ def setup(self) -> None: model_id, _ = _QWEN_LM_VARIANTS_INFO[self.model_variant] self.weight_file = str(Path(self.model_dir) / model_id) + vllm_kwargs = dict(self.vllm_kwargs) + if self.model_variant in _QWEN_LM_TEXT_ONLY_MULTIMODAL_VARIANTS: + vllm_kwargs.setdefault("limit_mm_per_prompt", {"image": 0, "video": 0}) self.llm = LLM( model=self.weight_file, quantization="fp8" if self.fp8 else None, - **self.vllm_kwargs, + **vllm_kwargs, ) self.sampling_params = SamplingParams( temperature=0.1, @@ -105,7 +111,12 @@ def setup(self) -> None: self.tokenizer = AutoTokenizer.from_pretrained(self.weight_file) def generate(self, inputs: list[dict[str, Any]]) -> list[str]: - formatted_inputs = self.tokenizer.apply_chat_template(inputs, tokenize=False, add_generation_prompt=True) + template_kwargs: dict[str, Any] = {} + if self.model_variant == "qwen3.5": + template_kwargs["enable_thinking"] = False + formatted_inputs = self.tokenizer.apply_chat_template( + inputs, tokenize=False, add_generation_prompt=True, **template_kwargs + ) results = self.llm.generate(formatted_inputs, sampling_params=self.sampling_params) return [result.outputs[0].text for result in results] diff --git a/nemo_curator/models/qwen_vl.py b/nemo_curator/models/qwen_vl.py index bc59d0a86f..938bbb654b 100644 --- a/nemo_curator/models/qwen_vl.py +++ b/nemo_curator/models/qwen_vl.py @@ -46,14 +46,19 @@ class SamplingParams: _QWEN3_VL_MODEL_ID = "Qwen/Qwen3-VL-8B-Instruct" _QWEN3_VL_MODEL_REVISION = "0c351dd" +_QWEN3_5_VL_MODEL_ID = "Qwen/Qwen3.5-9B" +_QWEN3_5_VL_MODEL_REVISION = "c202236" + _QWEN_VARIANTS_INFO: Final = { "qwen2.5": _QWEN2_5_VL_MODEL_ID, "qwen3": _QWEN3_VL_MODEL_ID, + "qwen3.5": _QWEN3_5_VL_MODEL_ID, } _QWEN_REVISION_INFO: Final = { "qwen2.5": _QWEN2_5_VL_MODEL_REVISION, "qwen3": _QWEN3_VL_MODEL_REVISION, + "qwen3.5": _QWEN3_5_VL_MODEL_REVISION, } _QWEN_VL_PIXEL_PARAMS: Final = { @@ -73,6 +78,14 @@ class SamplingParams: "video_max_pixels": 768 * 32 * 32, "video_total_pixels": 24576 * 32 * 32, }, + "qwen3.5": { + "image_factor": 32, + "min_pixels": 4 * 32 * 32, + "max_pixels": 16384 * 32 * 32, + "video_min_pixels": 128 * 32 * 32, + "video_max_pixels": 768 * 32 * 32, + "video_total_pixels": 24576 * 32 * 32, + }, } diff --git a/tests/models/test_prompt_formatter.py b/tests/models/test_prompt_formatter.py index 41395a59b3..097f7c6f87 100644 --- a/tests/models/test_prompt_formatter.py +++ b/tests/models/test_prompt_formatter.py @@ -31,6 +31,7 @@ def test_variant_mapping_contains_all_variants(self) -> None: expected_variants = { "qwen2.5", "qwen3", + "qwen3.5", "nemotron", "nemotron-bf16", "nemotron-fp8", @@ -43,6 +44,7 @@ def test_variant_mapping_qwen_hf_ids(self) -> None: """Test that Qwen variants have correct HuggingFace IDs.""" assert VARIANT_MAPPING["qwen2.5"] == "Qwen/Qwen2.5-VL-7B-Instruct" assert VARIANT_MAPPING["qwen3"] == "Qwen/Qwen3-VL-8B-Instruct" + assert VARIANT_MAPPING["qwen3.5"] == "Qwen/Qwen3.5-9B" def test_variant_mapping_nemotron_hf_ids(self) -> None: """Test that Nemotron variants have correct HuggingFace IDs.""" @@ -196,6 +198,63 @@ def test_create_qwen_message_special_characters(self) -> None: assert result[0]["content"][1]["text"] == special_prompt +class TestPromptFormatterQwen35: + """Test cases for PromptFormatter with the Qwen3.5 variant.""" + + @patch("nemo_curator.models.prompt_formatter.AutoProcessor") + def test_initialization_uses_qwen3_5_hf_id(self, mock_processor_class: Mock) -> None: + """Test that qwen3.5 loads the processor for its own checkpoint.""" + mock_processor_class.from_pretrained.return_value = Mock() + + formatter = PromptFormatter(prompt_variant="qwen3.5") + + assert formatter.prompt_variant == "qwen3.5" + mock_processor_class.from_pretrained.assert_called_once_with("Qwen/Qwen3.5-9B", trust_remote_code=True) + + @patch("nemo_curator.models.prompt_formatter.AutoProcessor") + def test_generate_inputs_routes_to_qwen_path(self, mock_processor_class: Mock) -> None: + """Test that qwen3.5 is dispatched to the Qwen formatter, not the Nemotron one.""" + mock_processor_instance = Mock() + mock_processor_class.from_pretrained.return_value = mock_processor_instance + mock_processor_instance.apply_chat_template.return_value = "formatted_prompt" + + formatter = PromptFormatter(prompt_variant="qwen3.5") + video_tensor = torch.randint(0, 255, (1, 3, 224, 224), dtype=torch.uint8) + + result = formatter.generate_inputs(prompt="Test prompt", video_inputs=video_tensor) + + assert result["prompt"] == "formatted_prompt" + # The Qwen message shape uses a structured content list; Nemotron uses a "