nemo_automodel.components.speculative.regenerate_vlm

Regenerate image-conversation answers with a vision-language target model.

ViSpec’s stage-2 corpus is built by throwing away the original (short) captions of an image-caption dataset and having the target VLM itself answer each prompt, with a length instruction appended so the answers are long enough to train a draft on. Two properties matter:

On-policy. The draft is supervised on the target’s own output distribution, so the training text has to come from the target.
Long assistant turns. Public multimodal SFT corpora answer in a sentence or two; a draft trained on those never sees the long-form decoding it has to accelerate. ViSpec appends “Please answer with at least 1000 words.” to every prompt to force long generations.

The output is ShareGPT-style JSONL plus a meta JSON, which nemo_automodel.components.datasets.vlm.datasets.make_meta_dataset reads directly, so the ViSpec recipe consumes it without a bespoke dataset class.

Example::

uv run python -m nemo_automodel.components.speculative.regenerate_vlm
—model Qwen/Qwen2.5-VL-7B-Instruct
—dataset liuhaotian/LLaVA-Pretrain
—image-root /data/LLaVA-Pretrain
—output-dir /data/vispec_stage2
—end 68000

Module Contents

Classes

Name	Description
`RegenerationConfig`	Settings for one regeneration run.

Functions

Name	Description
`_assert_prompt_round_trip`	Fail if the serialized row does not rebuild the exact generation prompt.
`_build_parser`	Build the command-line parser.
`_build_user_content`	Build the user-turn content parts, in the order the target sees them.
`_extract_prompt`	Pull the first human turn and image path out of a ShareGPT-style row.
`_generate_answer`	Run the target once and decode its answer.
`_load_source_dataset`	Load a source corpus from either a dataset ID or a local JSON file.
`_load_target_model`	Load the target VLM and place it on the available local device.
`_serialize_human_turn`	Serialize the generation prompt back into one ShareGPT `value` string.
`_write_meta`	Write the meta JSON that `make_meta_dataset` reads.
`main`	CLI entry point. Parses `argv` and returns the process exit code.
`regenerate`	Regenerate answers for a slice of the source corpus and write the shard.

Data

API

class nemo_automodel.components.speculative.regenerate_vlm.RegenerationConfig(
    model: str,
    dataset: str,
    split: str,
    image_root: str,
    output_dir: str,
    start: int,
    end: int,
    max_new_tokens: int,
    temperature: float,
    shuffle_seed: int,
    length_instruction: str,
    image_max_pixels: int | None = None,
    image_min_pixels: int | None = None
)

Dataclass

Settings for one regeneration run.

dataset

str

end

int

image_max_pixels

int | None = None

image_min_pixels

int | None = None

image_root

str

length_instruction

str

max_new_tokens

int

model

str

output_dir

str

shuffle_seed

int

split

str

start

int

temperature

float

nemo_automodel.components.speculative.regenerate_vlm._assert_prompt_round_trip(
    row: dict,
    expected_content: list[dict],
    image_root: str
) -> None

Fail if the serialized row does not rebuild the exact generation prompt.

The answer is only a valid supervision target for the prompt it was sampled under, so the prompt stage 2 reconstructs has to be the prompt this script generated from. That reconstruction happens in :func:convert_sharegpt_to_conversation, which is what make_meta_dataset runs over data.jsonl, so the check calls the real parser rather than a copy of its rules: a change to either side that breaks the correspondence surfaces here instead of silently training the draft on a prompt the target never answered.

Parameters:

row

dict

The row about to be written to data.jsonl.

expected_content

list[dict]

The user content that was handed to the target.

image_root

str

The media_dir recorded in meta.json.

Raises:

ValueError: If the rebuilt user turn differs from expected_content.

nemo_automodel.components.speculative.regenerate_vlm._build_parser() -> argparse.ArgumentParser

Build the command-line parser.

nemo_automodel.components.speculative.regenerate_vlm._build_user_content(
    prompt_text: str,
    image_path: str,
    length_instruction: str
) -> list[dict]

Build the user-turn content parts, in the order the target sees them.

The order (source text, image, length instruction) matches the reference implementation’s preprocess_function.

Parameters:

prompt_text

str

The source row’s human turn, image placeholder stripped.

image_path

str

Path to the image for this row.

length_instruction

str

Sentence appended after the image, or empty.

Returns: list[dict]

The content list of a single multimodal user turn.

nemo_automodel.components.speculative.regenerate_vlm._extract_prompt(
    example: dict
) -> tuple[str, str] | None

Pull the first human turn and image path out of a ShareGPT-style row.

Parameters:

example

dict

One source row, with a conversations list and an image path.

Returns: tuple[str, str] | None

(prompt_text, image_path), or None when the row has no usable

nemo_automodel.components.speculative.regenerate_vlm._generate_answer(
    model,
    processor,
    messages: list[dict],
    config: nemo_automodel.components.speculative.regenerate_vlm.RegenerationConfig
) -> str

Run the target once and decode its answer.

Parameters:

model

The loaded target VLM.

processor

The target’s processor.

messages

list[dict]

Chat messages wrapping :func:_build_user_content.

config

RegenerationConfig

The active regeneration settings.

Returns: str

The decoded assistant answer, without the prompt prefix.

nemo_automodel.components.speculative.regenerate_vlm._load_source_dataset(
    dataset_path: str,
    split: str
)

Load a source corpus from either a dataset ID or a local JSON file.

Parameters:

dataset_path

str

HuggingFace dataset identifier or local JSON/JSONL file.

split

str

Dataset split name passed to datasets.load_dataset.

Returns:

The loaded HuggingFace dataset split.

nemo_automodel.components.speculative.regenerate_vlm._load_target_model(
    model_path: str
) -> torch.nn.Module

Load the target VLM and place it on the available local device.

device_map="auto" requires the optional accelerate package. The regeneration utility only needs one local target replica, so explicit placement works in both minimal and Accelerate-enabled environments.

Parameters:

model_path

str

HuggingFace model identifier or local checkpoint directory.

Returns: torch.nn.Module

The evaluated target model on CUDA when available, otherwise CPU.

nemo_automodel.components.speculative.regenerate_vlm._serialize_human_turn(
    prompt_text: str,
    length_instruction: str
) -> str

Serialize the generation prompt back into one ShareGPT value string.

make_meta_dataset rebuilds the user turn by splitting this string on the <image> placeholder positionally, so writing the parts in the order they were generated is what makes the training prefix reproduce the generation prefix. :func:_assert_prompt_round_trip verifies that per row.

Parameters:

prompt_text

str

The source row’s human turn, image placeholder stripped.

length_instruction

str

Sentence appended after the image, or empty.

Returns: str

The value written to the human turn of data.jsonl.

nemo_automodel.components.speculative.regenerate_vlm._write_meta(
    output_dir: pathlib.Path,
    image_root: str
) -> pathlib.Path

Write the meta JSON that make_meta_dataset reads.

Parameters:

output_dir

Path

Directory holding the generated data.jsonl.

image_root

str

Directory the JSONL’s relative image paths resolve against.

Returns: Path

Path to the written meta file.

nemo_automodel.components.speculative.regenerate_vlm.main(
    argv: list[str] | None = None
) -> int

CLI entry point. Parses argv and returns the process exit code.

nemo_automodel.components.speculative.regenerate_vlm.regenerate(
    config: nemo_automodel.components.speculative.regenerate_vlm.RegenerationConfig
) -> pathlib.Path

Regenerate answers for a slice of the source corpus and write the shard.

Parameters:

config

RegenerationConfig

The regeneration settings.

Returns: Path

Path to the written data.jsonl.

nemo_automodel.components.speculative.regenerate_vlm.IMAGE_PLACEHOLDER = '<image>'

nemo_automodel.components.speculative.regenerate_vlm.LENGTH_INSTRUCTION = 'Please answer with at least 1000 words.'

nemo_automodel.components.speculative.regenerate_vlm.SHAREGPT_COLUMNS = {'messages': 'conversations', 'images': 'images'}

nemo_automodel.components.speculative.regenerate_vlm.SHAREGPT_TAGS = {'role_tag': 'from', 'content_tag': 'value', 'user_tag': 'human', 'assistant_tag...

nemo_automodel.components.speculative.regenerate_vlm.logger = logging.getLogger(__name__)