feat: integrate PreSel and Adapt-inf selectors

.gitattributes

@@ -0,0 +1,3 @@
+# Preserve the vendored function bodies in their algorithm implementations.
+src/dataflex/train/selector/adapt_inf_selector.py -whitespace
+src/dataflex/offline_selector/record_presel_selector.py -whitespace

examples/train_lora/selectors/adapt_inf.yaml

@@ -0,0 +1,58 @@
+# LLaVA-1.5 example. Override model/data/output paths.
+# 32 processes x batch4 x accumulation1 = global batch128.
+# DataFlex schedule example; 732 steps are training visits, not a unique-record budget.
+# CCS output is used unchanged, and previously used records may be selected again.
+# Install dataflex[multimodal-selection] for the author's TRAK projector.
+# Inspect the WSS grid and set num_clusters in components.yaml for your dataset.
+model_name_or_path: models/llava-v1.5-7b-stage1-hf-tokenizer-fixed
+stage: sft
+do_train: true
+finetuning_type: lora
+lora_target: q_proj,v_proj
+lora_rank: 128
+lora_alpha: 256
+dataset: llava620k
+dataset_dir: data
+template: llava
+cutoff_len: 2048
+packing: false
+preprocessing_num_workers: 16
+dataloader_num_workers: 4
+output_dir: ../dataflex_saves/adapt_inf_training
+report_to: none
+seed: 42
+per_device_train_batch_size: 4
+gradient_accumulation_steps: 1
+learning_rate: 0.0002
+num_train_epochs: 1.0
+bf16: true
+gradient_checkpointing: true
+use_reentrant_gc: false
+flash_attn: fa2
+save_strategy: steps
+save_steps: 244
+logging_steps: 1
+train_type: dynamic_select
+disable_tqdm: true
+components_cfg_file: src/dataflex/configs/components.yaml
+component_name: adapt_inf
+warmup_step: 100
+update_step: 382
+update_times: 10
+train_step: 732
+overwrite_output_dir: false
+full_determinism: true
+freeze_multi_modal_projector: false
+freeze_vision_tower: true
+additional_target: multi_modal_projector
+data_seed: 42
+save_total_limit: 3
+lr_scheduler_type: cosine
+warmup_ratio: 0.03
+weight_decay: 0.0
+ddp_timeout: 86400
+max_steps: 732
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+use_fast_tokenizer: false
+lora_dropout: 0.0

examples/train_lora/selectors/presel_reference.yaml

@@ -0,0 +1,48 @@
+# LLaVA-1.5 example. Override model/data/output paths.
+# 32 processes x batch4 x accumulation1 = global batch128.
+model_name_or_path: models/llava-v1.5-7b-stage1-hf-tokenizer-fixed
+stage: sft
+do_train: true
+finetuning_type: lora
+lora_target: q_proj,v_proj
+lora_rank: 128
+lora_alpha: 256
+dataset: presel_reference
+dataset_dir: outputs/presel/reference
+template: llava
+cutoff_len: 2048
+packing: false
+preprocessing_num_workers: 16
+dataloader_num_workers: 4
+output_dir: ../dataflex_saves/presel_reference
+report_to: none
+seed: 42
+per_device_train_batch_size: 4
+gradient_accumulation_steps: 1
+learning_rate: 0.0002
+num_train_epochs: 1.0
+bf16: true
+gradient_checkpointing: true
+use_reentrant_gc: false
+flash_attn: fa2
+save_strategy: steps
+save_steps: 244
+logging_steps: 1
+train_type: static
+disable_tqdm: true
+overwrite_output_dir: false
+full_determinism: true
+freeze_multi_modal_projector: false
+freeze_vision_tower: true
+additional_target: multi_modal_projector
+data_seed: 42
+save_total_limit: 3
+lr_scheduler_type: cosine
+warmup_ratio: 0.03
+weight_decay: 0.0
+ddp_timeout: 86400
+max_steps: -1
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+use_fast_tokenizer: false
+lora_dropout: 0.0

examples/train_lora/selectors/presel_static.yaml

@@ -0,0 +1,48 @@
+# LLaVA-1.5 example. Override model/data/output paths.
+# 32 processes x batch4 x accumulation1 = global batch128.
+model_name_or_path: models/llava-v1.5-7b-stage1-hf-tokenizer-fixed
+stage: sft
+do_train: true
+finetuning_type: lora
+lora_target: q_proj,v_proj
+lora_rank: 128
+lora_alpha: 256
+dataset: presel_selected
+dataset_dir: outputs/presel/selected
+template: llava
+cutoff_len: 2048
+packing: false
+preprocessing_num_workers: 16
+dataloader_num_workers: 4
+output_dir: ../dataflex_saves/presel_selected
+report_to: none
+seed: 42
+per_device_train_batch_size: 4
+gradient_accumulation_steps: 1
+learning_rate: 0.0002
+num_train_epochs: 1.0
+bf16: true
+gradient_checkpointing: true
+use_reentrant_gc: false
+flash_attn: fa2
+save_strategy: steps
+save_steps: 244
+logging_steps: 1
+train_type: static
+disable_tqdm: true
+overwrite_output_dir: false
+full_determinism: true
+freeze_multi_modal_projector: false
+freeze_vision_tower: true
+additional_target: multi_modal_projector
+data_seed: 42
+save_total_limit: 3
+lr_scheduler_type: cosine
+warmup_ratio: 0.03
+weight_decay: 0.0
+ddp_timeout: 86400
+max_steps: -1
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+use_fast_tokenizer: false
+lora_dropout: 0.0

pyproject.toml

@@ -41,8 +41,9 @@ torch = ["torch>=2.4.0,<=2.10.0", "torchvision>=0.19.0,<=0.21.0", "torchaudio>=2
 torch-npu = ["torch==2.4.0", "torch-npu==2.4.0.post2", "decorator"]
 # Needed by every examples/deepspeed config; 0.17+ cannot import on the torch 2.6 the README pins.
 deepspeed = ["deepspeed>=0.15.0,<0.17.0"]
-# Only the LESS selector needs TRAK, and it is a heavy build.
+# TRAK is used by LESS and Adapt-inf; install it only with those extras.
 less = ["traker"]
+multimodal-selection = ["scikit-learn>=1.3", "faiss-cpu>=1.8", "traker"]
 metrics = ["nltk", "jieba", "rouge-chinese"]
 flash-attn = ["flash-attn>=2.5.0"]
 liger-kernel = ["liger-kernel>=0.5.5"]

src/dataflex/cli.py

@@ -241,6 +241,11 @@ def main():
         # Only print version and welcome
         print_welcome()
         return
+    elif command == "presel":
+        from dataflex.offline_selector.offline_presel_selector import main as presel_main
+
+        presel_main(sys.argv[1:])
+        return
     elif command != 'train':
         raise ValueError(f'Unknown command: {command}')
     cfg = read_args()
@@ -285,4 +290,4 @@ def main():
         )
         sys.exit(process.returncode)
     else:
-        run_exp()
+        run_exp(cfg)

src/dataflex/configs/components.yaml

@@ -1,5 +1,17 @@
 # configs/components.yaml
 selectors:
+  adapt_inf:
+    name: adapt_inf
+    params:
+      cache_dir: ../dataflex_saves/adapt_inf
+      num_clusters: 10  # explicit choice; inspect the saved WSS grid for the actual dataset
+      cluster_grid: [5, 10, 15, 20, 25, 30, 35, 40, 45, 50]
+      proj_dim: 8192
+      projection_seed: 0  # author TRAK projection seed
+      parameter_pattern: '\.layers\.16\..*lora_'
+      seed: 42
+      histogram_bins: 50
+      dataset_revision: null
   # Dynamic selectors
   nice:
     name: nice
@@ -542,4 +554,4 @@ pipelines:
           component: loss
           mode: draw           # when used alone, it takes the budget for the entire round
           warmup_step: 10
-          every: 1
+          every: 1

src/dataflex/launcher.py

@@ -226,8 +226,8 @@ def launch():
         patch_lego_get_dataset()
 
     from llamafactory.train.tuner import run_exp
-    run_exp()
+    run_exp(cfg)
 
 
 if __name__ == "__main__":
-    launch()
+    launch()