Add helper function to get data module for llava

nemo/collections/vlm/api.py

@@ -0,0 +1,30 @@
+def get_llava_data_module(model_id: str, data_path: str, mbs: int, gbs: int):
+    from transformers import AutoProcessor
+
+    from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer
+    from nemo.collections.multimodal.data.energon import SimpleMultiModalDataModule
+    from nemo.collections.multimodal.data.energon.config import MultiModalSampleConfig
+    from nemo.collections.vlm import LlavaNextTaskEncoder
+
+    processor = AutoProcessor.from_pretrained(model_id)
+    tokenizer = AutoTokenizer(model_id)
+
+    multimodal_sample_config = MultiModalSampleConfig()
+    # Setting system prompt to empty string
+    multimodal_sample_config.conversation_template_config.system = ''
+
+    task_encoder = LlavaNextTaskEncoder(
+        tokenizer=tokenizer.tokenizer,
+        image_processor=processor.image_processor,
+        multimodal_sample_config=multimodal_sample_config,
+    )
+    return SimpleMultiModalDataModule(
+        path=data_path,
+        tokenizer=tokenizer,
+        image_processor=processor.image_processor,
+        num_workers=32,
+        micro_batch_size=mbs,
+        global_batch_size=gbs,
+        multimodal_sample_config=multimodal_sample_config,
+        task_encoder=task_encoder,
+    )

scripts/vlm/llava_next_nemo_run.py

@@ -15,6 +15,7 @@
 import nemo_run as run
 
 from nemo.collections import vlm
+from nemo.collections.vlm.api import get_llava_data_module
 
 
 def configure_recipe(nodes: int = 1, gpus_per_node: int = 8, pretrain=False, language_model_from_pretrained=None):
@@ -50,23 +51,26 @@ def local_executor_torchrun(nodes: int = 1, devices: int = 8) -> run.LocalExecut
     return executor
 
 
-def run_pretraining(language_model_from_pretrained=None):
+def run_finetuning():
     # pylint: disable=C0115,C0116
-    recipe = configure_recipe(pretrain=True, language_model_from_pretrained=language_model_from_pretrained)
+    recipe = configure_recipe(pretrain=False)
     executor = local_executor_torchrun(nodes=recipe.trainer.num_nodes, devices=recipe.trainer.devices)
 
     run.run(recipe, executor=executor)
 
 
-def run_finetuning():
+def run_pretraining():
     # pylint: disable=C0115,C0116
-    recipe = configure_recipe(pretrain=False)
+    recipe = configure_recipe(pretrain=True)
     executor = local_executor_torchrun(nodes=recipe.trainer.num_nodes, devices=recipe.trainer.devices)
+    print(f"recipe.model:{recipe.model}")
+    recipe.data = run.Config(
+        get_llava_data_module, model_id="llava-hf/llava-v1.6-vicuna-7b-hf", data_path="/data/path", mbs=2, gbs=8
+    )
 
     run.run(recipe, executor=executor)
 
 
 # This condition is necessary for the script to be compatible with Python's multiprocessing module.
 if __name__ == "__main__":
-    run_pretraining(language_model_from_pretrained='/root/.cache/nemo/models/lmsys/vicuna-7b-v1.5/')
-    # run_finetuning()
+    run_pretraining()