From ce0a6af840b01991585f21a2de5f2e114592c944 Mon Sep 17 00:00:00 2001 From: Ranjit Rajan Date: Wed, 1 Oct 2025 17:01:33 -0700 Subject: [PATCH 01/17] Add NemotronH VLM support for video captioning Signed-off-by: Ranjit Rajan --- nemo_curator/models/nemotron_h_vl.py | 225 ++ .../models/nemotron_prompt_formatter.py | 137 + nemo_curator/models/prompt_formatter.py | 50 +- .../video/caption/caption_generation.py | 41 +- .../video/caption/caption_preparation.py | 26 +- nemo_curator/tasks/video.py | 14 +- pyproject.toml | 3 +- .../video_split_clip_example.py | 27 +- uv.lock | 2428 ++++++++++------- 9 files changed, 1858 insertions(+), 1093 deletions(-) create mode 100644 nemo_curator/models/nemotron_h_vl.py create mode 100644 nemo_curator/models/nemotron_prompt_formatter.py diff --git a/nemo_curator/models/nemotron_h_vl.py b/nemo_curator/models/nemotron_h_vl.py new file mode 100644 index 0000000000..b587fa3150 --- /dev/null +++ b/nemo_curator/models/nemotron_h_vl.py @@ -0,0 +1,225 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +from pathlib import Path +from typing import Any + +from loguru import logger + +from nemo_curator.models.base import ModelInterface +from nemo_curator.utils import grouping +from nemo_curator.utils.hf_download_utils import download_model_from_hf + +# Constants for prompt processing +VIDEO_TAG_SPLIT_MAX = 1 +EXPECTED_VIDEO_TAG_PARTS = 2 + +try: + from vllm import LLM, SamplingParams + + VLLM_AVAILABLE = True +except ImportError: + VLLM_AVAILABLE = False + + class LLM: + pass + + class SamplingParams: + pass + + +# HuggingFace model IDs (will be updated when models are published) +_NEMOTRON_H_NANO_MODEL_ID = None +_NEMOTRON_H_NANO_MODEL_REVISION = None + + +class NemotronHVL(ModelInterface): + """NemotronH hybrid Mamba-Attention VLM for video captioning.""" + + def __init__( # noqa: PLR0913 + self, + model_dir: str, + model_variant: str = "nemotron", + caption_batch_size: int = 8, + max_output_tokens: int = 512, + stage2_prompt_text: str | None = None, + verbose: bool = False, + ): + self.model_dir = model_dir + self.model_variant = model_variant + self.caption_batch_size = caption_batch_size + self.max_output_tokens = max_output_tokens + self.stage2_prompt = stage2_prompt_text if stage2_prompt_text else "Please refine this caption: " + self.verbose = verbose + + if _NEMOTRON_H_NANO_MODEL_ID is not None: + self.weight_file = str(Path(model_dir) / _NEMOTRON_H_NANO_MODEL_ID) + else: + # Local checkpoint: model_dir is the checkpoint path itself + self.weight_file = str(Path(model_dir)) + + @property + def model_id_names(self) -> list[str]: + """Return model ID from config.json or HuggingFace ID.""" + if _NEMOTRON_H_NANO_MODEL_ID is not None: + return [_NEMOTRON_H_NANO_MODEL_ID] + + # Read from config.json if available + try: + config_path = Path(self.weight_file) / "config.json" + with open(config_path) as f: + config = json.load(f) + return [config.get("_name_or_path", Path(self.weight_file).name)] + except (FileNotFoundError, json.JSONDecodeError, KeyError): + return [Path(self.weight_file).name] + + def setup(self) -> None: + if not VLLM_AVAILABLE: + msg = "vllm is required for NemotronHVL but is not installed. Please install vllm: pip install vllm" + raise ImportError(msg) + + # Use V0 engine to avoid flashinfer issues for now + os.environ["VLLM_USE_V1"] = "0" + logger.info("Using vLLM V0 engine.") + + os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" + + self.model = LLM( + model=self.weight_file, + trust_remote_code=True, + tensor_parallel_size=1, + gpu_memory_utilization=0.9, + max_model_len=32768, + limit_mm_per_prompt={"video": 1}, + ) + + self.sampling_params = SamplingParams( + temperature=0.6, + max_tokens=self.max_output_tokens, + top_p=0.95, + stop=["", "<|endoftext|>", "", ""], + ) + + logger.info( + f"NemotronHVL initialized: variant={self.model_variant}, " + f"TP=1, GPU_util=0.9, max_len=32768" + ) + + def _refine_caption_prompt(self, original_prompt: str, refinement_text: str) -> str: + """Create a refined prompt for stage 2 captioning.""" + if "