add llava and instructblip

Former-commit-id: 142fb6f4541a1acfefe66ff2574dabde53b00c06
This commit is contained in:
BUAADreamer
2024-04-25 00:22:43 +08:00
parent 1451297c78
commit 12c51655ce
16 changed files with 273 additions and 214 deletions

View File

@@ -1,21 +1,14 @@
# Inspired by: https://github.com/huggingface/transformers/blob/v4.34.1/examples/pytorch/summarization/run_summarization.py
import os
from typing import TYPE_CHECKING, List, Optional
import torch
from PIL import Image
from torch.utils.data import Dataset
from transformers import DataCollatorForSeq2Seq, LlavaNextForConditionalGeneration, AutoModelForVision2Seq
from ...data import split_dataset, get_mm_dataset
from ...extras.constants import IGNORE_INDEX
from ...extras.misc import get_logits_processor
from ...extras.ploting import plot_loss
from ...model import load_model, load_tokenizer, load_processor, load_mm_model
from ...model import load_tokenizer, load_processor, load_mm_model
from ..utils import create_modelcard_and_push
from .metric import ComputeMetrics
from .trainer import CustomSeq2SeqTrainer
from .collator import DataCollatorForVis2Seq, ImageCaptioningDataset
from .collator import DataCollatorForVis2Seq
if TYPE_CHECKING:
from transformers import Seq2SeqTrainingArguments, TrainerCallback
@@ -32,28 +25,27 @@ def run_sft_mm(
callbacks: Optional[List["TrainerCallback"]] = None,
):
processor = load_processor(model_args)
tokenizer = processor.tokenizer
model = load_mm_model(processor, model_args, finetuning_args, training_args.do_train)
tokenizer = load_tokenizer(model_args)
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 %}"""
tokenizer.chat_template = CHAT_TEMPLATE
processor.tokenizer = tokenizer
use_clm = True
if "blip" in model_args.model_name_or_path:
use_clm = False
model = load_mm_model(processor, model_args, finetuning_args, training_args.do_train, use_clm=use_clm)
dataset = get_mm_dataset(processor, model_args, data_args, training_args, stage="sft")
if training_args.predict_with_generate:
tokenizer.padding_side = "left" # use left-padding in generation
if getattr(model, "is_quantized", False) and not training_args.do_train:
setattr(model, "_hf_peft_config_loaded", True) # hack here: make model compatible with prediction
splited_dataset = split_dataset(dataset, data_args, training_args)
splited_dataset['train_dataset'].set_format(type=splited_dataset['train_dataset'].format["type"],
columns=list(splited_dataset['train_dataset'].features.keys()))
splited_dataset['eval_dataset'].set_format(type=splited_dataset['eval_dataset'].format["type"],
columns=list(splited_dataset['eval_dataset'].features.keys()))
train_dataset = ImageCaptioningDataset(splited_dataset['train_dataset'], data_args.image_path, processor)
eval_dataset = ImageCaptioningDataset(splited_dataset['eval_dataset'], data_args.image_path, processor)
train_dataset = dataset
eval_dataset = dataset
data_collator = DataCollatorForVis2Seq(
processor=processor,
use_qformer=model_args.use_qformer,
)
# Override the decoding parameters of Seq2SeqTrainer
training_args.generation_max_length = training_args.generation_max_length or data_args.cutoff_len
training_args.generation_num_beams = data_args.eval_num_beams or training_args.generation_num_beams
training_args.remove_unused_columns = False
# Initialize our Trainer
trainer = CustomSeq2SeqTrainer(
@@ -67,7 +59,6 @@ def run_sft_mm(
train_dataset=train_dataset,
eval_dataset=eval_dataset,
)
# Keyword arguments for `model.generate`
gen_kwargs = generating_args.to_dict()
gen_kwargs["eos_token_id"] = [tokenizer.eos_token_id] + tokenizer.additional_special_tokens_ids