sft_utils

sft_utils#

The supervised fine-tuning core: dataset loading and formatting, chain-of-thought prompt filling, LoRA and full-parameter trainer construction, and the temperature-based multi-task sampler.

medvision_bm.sft.sft_utils.is_main_process()[source]#
medvision_bm.sft.sft_utils.safe_print(*args, force=False, **kwargs)[source]#

Print only on main process unless force=True.

medvision_bm.sft.sft_utils.broadcast_int_from_main(value, src=0)[source]#
medvision_bm.sft.sft_utils.get_cgroup_limited_cpus()[source]#
medvision_bm.sft.sft_utils.get_image_info_for_medvision_dataset(doc)[source]#

Get the image modality and label name for a MedVision sample.

The sample’s taskType (defined in MedVision.py on the dataset repo) selects the dataset-specific preprocessing module, from which the image modality and the human-readable label name are looked up.

Parameters:

doc – A data sample from the MedVision dataset. Must contain taskType, dataset_name and taskID; label is present for all task types except the Biometrics-From-Landmarks (angle / distance) tasks.

Returns:

(image_modality, label_name), where label_name is None for tasks that have no label (the angle / distance tasks).

Return type:

tuple

Raises:

ValueError – If taskType is not one of the valid task types, or the dataset is not registered in DATASETS_NAME2PACKAGE.

medvision_bm.sft.sft_utils.normalize_ct_img(img, window_width, window_level)[source]#

Normalizes CT Hounsfield Units to [0, 255] based on W and L.

medvision_bm.sft.sft_utils.normalize_general_img(img)[source]#

Standard min-max normalization to [0, 255] for MR, PET, etc.

medvision_bm.sft.sft_utils.normalize_img(doc, img_2d)[source]#

Convert document to image with scale bar added.

medvision_bm.sft.sft_utils.img_proccessor_nii2png_save2disk(example, new_shape_hw=None)[source]#
medvision_bm.sft.sft_utils.img_proccessor_nii2png_save2dataset(example, new_shape_hw=None)[source]#
medvision_bm.sft.sft_utils.safe_concat_align_top_keys(datasets_list, fill_value=None)[source]#

Concatenate Hugging Face datasets even if they have different top-level keys. Missing columns are added and filled with fill_value.

medvision_bm.sft.sft_utils.safe_concat_align_dict_keys(datasets_list, dict_cols=None, fill_value=None)[source]#

Concatenate Hugging Face datasets even if dict columns have different keys.

Parameters:

  • datasets_list (list[Dataset]) – list of datasets to concatenate

  • dict_cols (list[str] | None) – names of columns containing dicts (if None, auto-detects)

  • fill_value (any) – value used to fill missing keys (default: None)

Returns:

concatenated dataset

Return type:

Dataset

medvision_bm.sft.sft_utils.safe_concatenate_datasets(datasets_list)[source]#
medvision_bm.sft.sft_utils.group_train_test_split(dataset, group_column, test_size, seed=None, stratify_column=None)[source]#

Splits a HF Dataset into train and validation sets ensuring samples with the same value in ‘group_column’ are in the same split.

Parameters:

  • dataset – The HF Dataset to split.

  • group_column – The column name to group by (e.g., ‘image_file’).

  • test_size – If float < 1.0, represents fraction of samples to aim for. If int >= 1, represents exact number of samples to aim for.

  • seed – Random seed for shuffling.

  • stratify_column – Optional column name to stratify by (e.g., ‘dataset_name’). A stratum is the set of volumes belonging to one unique value of this column (e.g., all volumes from ‘BraTS24’ form one stratum, all volumes from ‘AMOS22’ form another). When provided, volumes are interleaved round-robin across strata — one volume per stratum per round — before the greedy allocation loop runs. This guarantees every stratum contributes at least one volume to val before any stratum gets a second, preventing a few large-volume datasets from monopolising the val quota.

Returns:

DatasetDict containing ‘train’ and ‘validation’.

medvision_bm.sft.sft_utils.load_split_limit_dataset(tasks_list_json_path, limit_train_sample, limit_val_sample, num_workers_concat_datasets=4, tag_ds=None, download_mode='reuse_dataset_if_exists')[source]#

Load MedVision tasks, concatenate them, and split into train/validation.

Reads the task list from tasks_list_json_path and loads each task’s _Train split in parallel (falling back to single-threaded loading when any dataset was just downloaded, to avoid cache conflicts). The per-task datasets are concatenated in the deterministic JSON order (not arrival order) so the seeded shuffle and split downstream stay reproducible, then split into train and validation sets grouped by image_file (to prevent 3D-volume leakage) and stratified by dataset_name.

Parameters:

  • tasks_list_json_path (str) – Path to the JSON file whose keys are the task names to load.

  • limit_train_sample (int) – Cap on training samples after concatenation. Use a value < 0 for no limit or > 0 for a fixed cap; 0 is rejected.

  • limit_val_sample (int) – Target size of the validation split; must be > 0.

  • num_workers_concat_datasets (int) – Requested worker processes for parallel loading; clamped to the CPU count and task count. Defaults to 4.

  • tag_ds (str) – Tag embedded in task names (<dataset_name>_<tag_ds>), used to recover the dataset name from each task. Required.

  • download_mode (str) – Passed through to the dataset loader. Defaults to "reuse_dataset_if_exists".

Returns:

A dict with "train" and "validation" splits.

Return type:

DatasetDict

Raises:

  • AssertionError – If limit_val_sample is not > 0, limit_train_sample is 0, tag_ds is None, or MedVision_DATA_DIR is unset.

  • RuntimeError – If any task fails to load.

medvision_bm.sft.sft_utils.format_dataset(dataset, mapping_func, mapping_func_args, num_workers_format_dataset, writer_batch_size=1000)[source]#

Apply a formatting map function to a dataset with bounded parallelism.

Runs dataset.map(mapping_func, ...) to convert raw MedVision rows into the chat messages format expected by the SFT trainer. The number of worker processes is capped at the cgroup-limited CPU count.

Parameters:

  • dataset – A HuggingFace Dataset or DatasetDict to format.

  • mapping_func – The per-example formatting function passed to .map().

  • mapping_func_args (dict) – Keyword arguments forwarded to mapping_func via fn_kwargs.

  • num_workers_format_dataset (int) – Requested number of worker processes; clamped to the available CPU count.

  • writer_batch_size (int) – Number of rows buffered before writing to the Arrow cache. Defaults to 1000.

Returns:

The formatted dataset with the mapping function applied.

medvision_bm.sft.sft_utils.clean_dataset(dataset, keys_to_keep)[source]#

Drop all columns from a dataset except a whitelist of keys.

Maps over the dataset and deletes every key not present in keys_to_keep, keeping the cached rows small before training.

Parameters:

  • dataset – A HuggingFace Dataset or DatasetDict to prune.

  • keys_to_keep (list[str]) – Column names to retain; all other columns are removed.

Returns:

The dataset containing only the whitelisted columns.

medvision_bm.sft.sft_utils.prepare_dataset(*, tasks_list_json_path, limit_train_sample, limit_val_sample, mapping_func, model_family_name, base_model_hf, num_workers_concat_datasets=4, num_workers_format_dataset=32, tag_ds=None, process_img=False, save_processed_img_to_disk=False, new_shape_hw=None, download_mode='reuse_dataset_if_exists')[source]#

Load, format, and prune a MedVision dataset for SFT in one call.

Combines load_split_limit_dataset(), format_dataset(), and clean_dataset(): it loads and splits the tasks, maps each example into the chat messages format via mapping_func, then keeps only the columns needed for training.

Parameters:

  • tasks_list_json_path (str) – Path to the JSON file listing the tasks.

  • limit_train_sample (int) – Training-sample cap (< 0 = no limit, > 0 = cap).

  • limit_val_sample (int) – Target validation-split size; must be > 0.

  • mapping_func – Per-example formatting function applied during mapping.

  • model_family_name (str) – Model family name passed to mapping_func as model_name (used for image-resize logic).

  • base_model_hf (str) – HuggingFace model id passed to mapping_func as model_hf.

  • num_workers_concat_datasets (int) – Worker processes for loading. Defaults to 4.

  • num_workers_format_dataset (int) – Worker processes for formatting. Defaults to 32.

  • tag_ds (str) – Tag embedded in task names; required by the loader.

  • process_img (bool) – If True, embed processed PNG images in the dataset (processed_images column). Not recommended (large cache).

  • save_processed_img_to_disk (bool) – If True, write processed PNGs to disk and store their paths in image_file_png.

  • new_shape_hw (tuple[int, int] | None) – Optional explicit (height, width) to resize slices to before formatting.

  • download_mode (str) – Passed through to the loader. Defaults to "reuse_dataset_if_exists".

Returns:

Train/validation splits containing only the retained columns (messages, labels, image_file, slice_dim, slice_idx, plus processed_images and/or image_file_png when enabled).

Return type:

DatasetDict

medvision_bm.sft.sft_utils.recompute_total_max_steps(trainer)[source]#

Recompute total planned update steps based on global dataset size, world size and desired epochs.

medvision_bm.sft.sft_utils.prepare_trainer(*, run_name, base_model_hf, lora_checkpoint_dir, data, make_collate_fn, per_device_train_batch_size=14, per_device_eval_batch_size=14, gradient_accumulation_steps=6, use_flash_attention_2=True, num_train_epochs=1, save_steps=100, eval_steps=50, logging_steps=50, save_total_limit=10, dataloader_num_workers=8, gradient_checkpointing=False, dataloader_pin_memory=True, push_LoRA=False, enable_temperature_sampler=False, temperature_sampler_T=3.0, temperature_sampler_task_column='__task_name', temperature_sampler_num_samples=-1)[source]#

Build a QLoRA SFTTrainer for MedVision SFT.

Loads base_model_hf in 4-bit NF4 quantization, attaches a LoRA adapter (all-linear target modules, with lm_head and embed_tokens also trained), and wraps it in an SFTTrainer configured for BF16 training and Weights & Biases logging. When enable_temperature_sampler is set, a temperature-weighted multi-task sampler is used instead of uniform sampling.

Parameters:

  • run_name (str) – Run name for logging / W&B.

  • base_model_hf (str) – HuggingFace id of the base image-text-to-text model.

  • lora_checkpoint_dir (str) – Output directory for adapter checkpoints.

  • data – DatasetDict with "train" and "validation" splits.

  • make_collate_fn – Factory called as make_collate_fn(processor) to build the data collator.

  • per_device_train_batch_size (int) – Per-device train batch size. Defaults to 14.

  • per_device_eval_batch_size (int) – Per-device eval batch size. Defaults to 14.

  • gradient_accumulation_steps (int) – Gradient accumulation steps. Defaults to 6.

  • use_flash_attention_2 (bool) – Use FlashAttention-2 when True, else eager attention. Defaults to True.

  • num_train_epochs (int) – Number of training epochs. Defaults to 1.

  • save_steps (int) – Steps between checkpoint saves. Defaults to 100.

  • eval_steps (int) – Steps between evaluations. Defaults to 50.

  • logging_steps (int) – Steps between log entries. Defaults to 50.

  • save_total_limit (int) – Max checkpoints to retain. Defaults to 10.

  • dataloader_num_workers (int) – DataLoader worker processes. Defaults to 8.

  • gradient_checkpointing (bool) – Enable gradient checkpointing. Defaults to False.

  • dataloader_pin_memory (bool) – Pin DataLoader memory. Defaults to True.

  • push_LoRA (bool) – Push the adapter to the Hub (private) when True. Defaults to False.

  • enable_temperature_sampler (bool) – Enable temperature-based multi-task sampling. Defaults to False.

  • temperature_sampler_T (float) – Sampling temperature. Defaults to 3.0.

  • temperature_sampler_task_column (str) – Column holding task labels. Defaults to "__task_name".

  • temperature_sampler_num_samples (int) – Draws per epoch; <= 0 keeps the training-set length. Defaults to -1.

Returns:

The configured trainer (a temperature-sampling subclass when enabled with more than one task).

Return type:

SFTTrainer

Raises:

ValueError – If the GPU does not support bfloat16.

medvision_bm.sft.sft_utils.prepare_trainer_fullFT(*, run_name, base_model_hf, checkpoint_dir, data, make_collate_fn, per_device_train_batch_size=1, per_device_eval_batch_size=1, gradient_accumulation_steps=16, use_flash_attention_2=True, num_train_epochs=1, save_steps=100, eval_steps=50, logging_steps=50, save_total_limit=5, dataloader_num_workers=4, gradient_checkpointing=True, dataloader_pin_memory=True, push_model=False, enable_temperature_sampler=False, temperature_sampler_T=3.0, temperature_sampler_task_column='__task_name', temperature_sampler_num_samples=-1)[source]#

Prepare an SFTTrainer for full parameter finetuning (no LoRA, no quantization).

Loads the model in BF16 without any PEFT adapter. All parameters are trained. Use a lower learning rate and cosine scheduler compared to the LoRA variant. When enable_temperature_sampler is set, a temperature-weighted multi-task sampler is used instead of uniform sampling.

Parameters:

  • run_name (str) – Run name for logging / W&B.

  • base_model_hf (str) – HuggingFace id of the base image-text-to-text model.

  • checkpoint_dir (str) – Output directory for checkpoints.

  • data – DatasetDict with "train" and "validation" splits.

  • make_collate_fn – Factory called as make_collate_fn(processor) to build the data collator.

  • per_device_train_batch_size (int) – Per-device train batch size. Defaults to 1.

  • per_device_eval_batch_size (int) – Per-device eval batch size. Defaults to 1.

  • gradient_accumulation_steps (int) – Gradient accumulation steps. Defaults to 16.

  • use_flash_attention_2 (bool) – Use FlashAttention-2 when True, else eager attention. Defaults to True.

  • num_train_epochs (int) – Number of training epochs. Defaults to 1.

  • save_steps (int) – Steps between checkpoint saves. Defaults to 100.

  • eval_steps (int) – Steps between evaluations. Defaults to 50.

  • logging_steps (int) – Steps between log entries. Defaults to 50.

  • save_total_limit (int) – Max checkpoints to retain. Defaults to 5.

  • dataloader_num_workers (int) – DataLoader worker processes. Defaults to 4.

  • gradient_checkpointing (bool) – Enable gradient checkpointing (on by default; required at 7B+ scale). Defaults to True.

  • dataloader_pin_memory (bool) – Pin DataLoader memory. Defaults to True.

  • push_model (bool) – Push the trained model to the Hub (private) when True. Defaults to False.

  • enable_temperature_sampler (bool) – Enable temperature-based multi-task sampling. Defaults to False.

  • temperature_sampler_T (float) – Sampling temperature. Defaults to 3.0.

  • temperature_sampler_task_column (str) – Column holding task labels. Defaults to "__task_name".

  • temperature_sampler_num_samples (int) – Draws per epoch; <= 0 keeps the training-set length. Defaults to -1.

Returns:

The configured trainer (a temperature-sampling subclass when enabled with more than one task).

Return type:

SFTTrainer

Raises:

ValueError – If the GPU does not support bfloat16.

medvision_bm.sft.sft_utils.merge_models(base_model_hf, lora_checkpoint_dir, merged_model_hf, merged_model_dir, push_to_hub)[source]#

Merge a LoRA adapter into its base model and optionally save/push it.

Loads the base model on CPU in fp32 (so the sub-BF16 LoRA delta is representable), merges the adapter with safe_merge=True, then optionally saves the merged model locally and/or pushes it to the Hugging Face Hub. This function is intended to be called only on the main process.

Parameters:

  • base_model_hf (str) – HuggingFace id of the base model.

  • lora_checkpoint_dir (str) – Directory of the trained LoRA adapter (also the source of the processor).

  • merged_model_hf (str) – Target Hub repo id for the merged model. Required only when push_to_hub is True.

  • merged_model_dir (str | None) – Local directory to save the merged model to; skipped when None.

  • push_to_hub (bool) – If True, push the merged model and processor to the Hub as a private repo.

Raises:

ValueError – If push_to_hub is True but merged_model_hf is None.

medvision_bm.sft.sft_utils.train_resume_from_checkpoint(trainer, last_checkpoint)[source]#
medvision_bm.sft.sft_utils.parse_args_multiTask()[source]#

Parse command-line arguments for SFT on the MedVision dataset.

medvision_bm.sft.sft_utils.check_model_supported(model_name)[source]#
medvision_bm.sft.sft_utils.parse_validate_args_multiTask()[source]#
medvision_bm.sft.sft_utils.parse_sample_limits(**kwargs)[source]#

Determine sample limits for each task with fallbacks.

Logic:

  • If task-specific limit > 0: use it

  • Else: use per-task limit

  • If task JSON path is None: set limit to 0 (task not used)

Returns:

(train_limit_AD, val_limit_AD,

train_limit_detect, val_limit_detect, train_limit_TL, val_limit_TL, train_limit_total)

Return type:

A tuple of sample limits

medvision_bm.sft.sft_utils.mask_non_assistant_turns(input_ids, labels, tokenizer)[source]#

Mask everything except assistant response content + its closing <|im_end|>.

Completion-only masking: for every assistant turn the header <|im_start|>assistant\n is masked (it is chat-template scaffolding the model never needs to generate), along with all system/user/tool turns; loss is computed only on the response tokens and the <|im_end|> that terminates the assistant turn. The trailing newline after <|im_end|> is also masked.