GohioAC/lava_next_vsft.py Secret

## lava_next_vsft.py
"""
python vsft.py \
    --dataset_name="HuggingFaceH4/llava-instruct-mix-vsft" \
    --model_name_or_path="llava-hf/llava-v1.6-mistral-7b-hf" \
    --report_to="tensorboard" \
    --learning_rate=2e-5 \
    --lr_scheduler_type="cosine" \
    --per_device_train_batch_size=8 \
    --gradient_accumulation_steps=1 \
    --output_dir="data/vsft-llava-1.5-7b-hf" \
    --logging_steps=1 \
    --num_train_epochs=1 \
    --gradient_checkpointing \
    --remove_unused_columns=False \
    --torch_dtype=float16 \
    --fp16=True \
    --max_seq_length=4096 \
    --attn_implementation="flash_attention_2"
"""

from contextlib import nullcontext
from trl.commands.cli_utils import SFTScriptArguments, TrlParser

import torch
from datasets import load_dataset

from tqdm.rich import tqdm
from transformers import AutoTokenizer, AutoProcessor, LlavaNextForConditionalGeneration

from trl import (
    ModelConfig,
    SFTConfig,
    SFTTrainer,
    get_peft_config,
    get_quantization_config,
    get_kbit_device_map,
)

tqdm.pandas()

if __name__ == "__main__":
    parser = TrlParser((SFTScriptArguments, SFTConfig, ModelConfig))
    sft_script_args, training_args, model_config = parser.parse_args_and_config()
    training_args.gradient_checkpointing_kwargs = dict(use_reentrant=False)

    LLAVA_CHAT_TEMPLATE = """{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. {% for message in messages %}{% if message['role'] == 'user' %}USER: {% else %}ASSISTANT: {% endif %}{% for item in message['content'] %}{% if item['type'] == 'text' %}{{ item['text'] }}{% elif item['type'] == 'image' %}<image>{% endif %}{% endfor %}{% if message['role'] == 'user' %} {% else %}{{eos_token}}{% endif %}{% endfor %}{% if add_generation_prompt %}ASSISTANT: {% endif %}"""

    torch_dtype = (
        model_config.torch_dtype
        if model_config.torch_dtype in ["auto", None]
        else getattr(torch, model_config.torch_dtype)
    )
    quantization_config = get_quantization_config(model_config)
    model_kwargs = dict(
        revision=model_config.model_revision,
        trust_remote_code=model_config.trust_remote_code,
        attn_implementation=model_config.attn_implementation,
        torch_dtype=torch_dtype,
        device_map=get_kbit_device_map() if quantization_config is not None else None,
        quantization_config=quantization_config,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_config.model_name_or_path, use_fast=True, padding_side="right"
    )
    tokenizer.chat_template = LLAVA_CHAT_TEMPLATE
    processor = AutoProcessor.from_pretrained(model_config.model_name_or_path)
    processor.tokenizer = tokenizer

    model = LlavaNextForConditionalGeneration.from_pretrained(
        model_config.model_name_or_path, **model_kwargs
    )

    class LLavaDataCollator:
        def __init__(self, processor):
            self.processor = processor

        def __call__(self, examples):
            texts = []
            images = []
            for example in examples:
                if len(example["images"]) > 1:
                    raise ValueError(
                        "This collator only supports one image per example"
                    )
                messages = example["messages"]
                text = self.processor.tokenizer.apply_chat_template(
                    messages, tokenize=False, add_generation_prompt=False
                )
                texts.append(text)
                images.append(example["images"][0])

            batch = self.processor(texts, images, return_tensors="pt", padding=True)

            labels = batch["input_ids"].clone()
            batch["labels"] = labels

            return batch

    data_collator = LLavaDataCollator(processor)

    raw_datasets = load_dataset(sft_script_args.dataset_name)
    train_dataset = raw_datasets[sft_script_args.dataset_train_split]
    eval_dataset = raw_datasets[sft_script_args.dataset_test_split]

    init_context = nullcontext()
    save_context = nullcontext()

    with init_context:
        trainer = SFTTrainer(
            model=model,
            args=training_args,
            train_dataset=train_dataset,
            eval_dataset=eval_dataset,
            dataset_text_field="text",  # need a dummy field
            tokenizer=tokenizer,
            peft_config=get_peft_config(model_config),
            callbacks=None,
            data_collator=data_collator,
            dataset_kwargs={"skip_prepare_dataset": True},
        )

    trainer.train()

    with save_context:
        trainer.save_model(training_args.output_dir)
	"""
	python vsft.py \
	--dataset_name="HuggingFaceH4/llava-instruct-mix-vsft" \
	--model_name_or_path="llava-hf/llava-v1.6-mistral-7b-hf" \
	--report_to="tensorboard" \
	--learning_rate=2e-5 \
	--lr_scheduler_type="cosine" \
	--per_device_train_batch_size=8 \
	--gradient_accumulation_steps=1 \
	--output_dir="data/vsft-llava-1.5-7b-hf" \
	--logging_steps=1 \
	--num_train_epochs=1 \
	--gradient_checkpointing \
	--remove_unused_columns=False \
	--torch_dtype=float16 \
	--fp16=True \
	--max_seq_length=4096 \
	--attn_implementation="flash_attention_2"
	"""

	from contextlib import nullcontext
	from trl.commands.cli_utils import SFTScriptArguments, TrlParser

	import torch
	from datasets import load_dataset

	from tqdm.rich import tqdm
	from transformers import AutoTokenizer, AutoProcessor, LlavaNextForConditionalGeneration

	from trl import (
	ModelConfig,
	SFTConfig,
	SFTTrainer,
	get_peft_config,
	get_quantization_config,
	get_kbit_device_map,
	)

	tqdm.pandas()

	if __name__ == "__main__":
	parser = TrlParser((SFTScriptArguments, SFTConfig, ModelConfig))
	sft_script_args, training_args, model_config = parser.parse_args_and_config()
	training_args.gradient_checkpointing_kwargs = dict(use_reentrant=False)

	LLAVA_CHAT_TEMPLATE = """{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. {% for message in messages %}{% if message['role'] == 'user' %}USER: {% else %}ASSISTANT: {% endif %}{% for item in message['content'] %}{% if item['type'] == 'text' %}{{ item['text'] }}{% elif item['type'] == 'image' %}<image>{% endif %}{% endfor %}{% if message['role'] == 'user' %} {% else %}{{eos_token}}{% endif %}{% endfor %}{% if add_generation_prompt %}ASSISTANT: {% endif %}"""

	torch_dtype = (
	model_config.torch_dtype
	if model_config.torch_dtype in ["auto", None]
	else getattr(torch, model_config.torch_dtype)
	)
	quantization_config = get_quantization_config(model_config)
	model_kwargs = dict(
	revision=model_config.model_revision,
	trust_remote_code=model_config.trust_remote_code,
	attn_implementation=model_config.attn_implementation,
	torch_dtype=torch_dtype,
	device_map=get_kbit_device_map() if quantization_config is not None else None,
	quantization_config=quantization_config,
	)
	tokenizer = AutoTokenizer.from_pretrained(
	model_config.model_name_or_path, use_fast=True, padding_side="right"
	)
	tokenizer.chat_template = LLAVA_CHAT_TEMPLATE
	processor = AutoProcessor.from_pretrained(model_config.model_name_or_path)
	processor.tokenizer = tokenizer

	model = LlavaNextForConditionalGeneration.from_pretrained(
	model_config.model_name_or_path, **model_kwargs
	)

	class LLavaDataCollator:
	def __init__(self, processor):
	self.processor = processor

	def __call__(self, examples):
	texts = []
	images = []
	for example in examples:
	if len(example["images"]) > 1:
	raise ValueError(
	"This collator only supports one image per example"
	)
	messages = example["messages"]
	text = self.processor.tokenizer.apply_chat_template(
	messages, tokenize=False, add_generation_prompt=False
	)
	texts.append(text)
	images.append(example["images"][0])

	batch = self.processor(texts, images, return_tensors="pt", padding=True)

	labels = batch["input_ids"].clone()
	batch["labels"] = labels

	return batch

	data_collator = LLavaDataCollator(processor)

	raw_datasets = load_dataset(sft_script_args.dataset_name)
	train_dataset = raw_datasets[sft_script_args.dataset_train_split]
	eval_dataset = raw_datasets[sft_script_args.dataset_test_split]

	init_context = nullcontext()
	save_context = nullcontext()

	with init_context:
	trainer = SFTTrainer(
	model=model,
	args=training_args,
	train_dataset=train_dataset,
	eval_dataset=eval_dataset,
	dataset_text_field="text", # need a dummy field
	tokenizer=tokenizer,
	peft_config=get_peft_config(model_config),
	callbacks=None,
	data_collator=data_collator,
	dataset_kwargs={"skip_prepare_dataset": True},
	)

	trainer.train()

	with save_context:
	trainer.save_model(training_args.output_dir)