add multimodal LLM BLIP-2 and InstructBLIP

Former-commit-id: a730f89a972f1a9d37c718c716f199cb8d4903b2

src/llmtuner/train/sftmm/__init__.py

@@ -0,0 +1,3 @@
+from .workflow import run_sft_mm
+
+__all__ = ["run_sft_mm"]

src/llmtuner/train/sftmm/collator.py

@@ -0,0 +1,69 @@
+import json
+import os
+from dataclasses import dataclass
+
+import torch
+from torch.utils.data import Dataset as Dataset_torch
+from datasets import Dataset
+from PIL import Image
+from transformers import AutoProcessor
+
+
+class ImageCaptioningDataset(Dataset_torch):
+    def __init__(self, dataset: Dataset, image_path: str, processor: AutoProcessor):
+        self.processor = processor
+        self.dataset = dataset
+        self.image_path = image_path
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        source = self.dataset[idx]
+        image_id = source['image']
+        image = Image.open(os.path.join(self.image_path, image_id))
+        convs = source['conversations']
+        prompt = convs[0]['value']
+        label = convs[1]['value']
+        image_inputs = self.processor(image, return_tensors="pt")
+        image_inputs = {k: v.squeeze() for k, v in image_inputs.items()}
+        inputs = {
+            "input_ids": prompt,
+            "labels": label,
+        }
+        for key in image_inputs:
+            inputs[key] = image_inputs[key]
+        return inputs
+
+
+@dataclass
+class DataCollatorForVis2Seq:
+    processor: AutoProcessor
+    use_qformer: bool = False
+
+    def __call__(self, features, return_tensors=None):
+        processed_batch = {}
+        for key in features[0].keys():
+            if key == 'pixel_values':
+                processed_batch[key] = torch.stack([example[key] for example in features])
+            elif key == 'input_ids':
+                text_inputs = self.processor.tokenizer(
+                    [example[key] for example in features], padding="max_length", return_tensors="pt",
+                    max_length=512,
+                )
+                processed_batch["input_ids"] = text_inputs["input_ids"]
+                processed_batch["attention_mask"] = text_inputs["attention_mask"]
+                if self.use_qformer:
+                    qformer_text_inputs = self.processor.qformer_tokenizer(
+                        [example[key] for example in features], padding="max_length", return_tensors="pt",
+                        max_length=512,
+                    )
+                    processed_batch["qformer_input_ids"] = qformer_text_inputs["input_ids"]
+                    processed_batch["qformer_attention_mask"] = qformer_text_inputs["attention_mask"]
+            elif key == 'labels':
+                text_inputs = self.processor.tokenizer(
+                    [example[key] for example in features], padding="max_length", return_tensors="pt",
+                    max_length=512,
+                )
+                processed_batch["labels"] = text_inputs["input_ids"]
+        return processed_batch

src/llmtuner/train/sftmm/metric.py

@@ -0,0 +1,61 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Dict, Sequence, Tuple, Union
+
+import numpy as np
+
+from ...extras.constants import IGNORE_INDEX
+from ...extras.packages import is_jieba_available, is_nltk_available, is_rouge_available
+
+
+if TYPE_CHECKING:
+    from transformers.tokenization_utils import PreTrainedTokenizer
+
+if is_jieba_available():
+    import jieba  # type: ignore
+
+if is_nltk_available():
+    from nltk.translate.bleu_score import SmoothingFunction, sentence_bleu
+
+if is_rouge_available():
+    from rouge_chinese import Rouge
+
+
+@dataclass
+class ComputeMetrics:
+    r"""
+    Wraps the tokenizer into metric functions, used in Seq2SeqPeftTrainer.
+    """
+
+    tokenizer: "PreTrainedTokenizer"
+
+    def __call__(self, eval_preds: Sequence[Union[np.ndarray, Tuple[np.ndarray]]]) -> Dict[str, float]:
+        r"""
+        Uses the model predictions to compute metrics.
+        """
+        preds, labels = eval_preds
+        score_dict = {"rouge-1": [], "rouge-2": [], "rouge-l": [], "bleu-4": []}
+
+        preds = np.where(preds != IGNORE_INDEX, preds, self.tokenizer.pad_token_id)
+        labels = np.where(labels != IGNORE_INDEX, labels, self.tokenizer.pad_token_id)
+
+        decoded_preds = self.tokenizer.batch_decode(preds, skip_special_tokens=True)
+        decoded_labels = self.tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+        for pred, label in zip(decoded_preds, decoded_labels):
+            hypothesis = list(jieba.cut(pred))
+            reference = list(jieba.cut(label))
+
+            if len(" ".join(hypothesis).split()) == 0 or len(" ".join(reference).split()) == 0:
+                result = {"rouge-1": {"f": 0.0}, "rouge-2": {"f": 0.0}, "rouge-l": {"f": 0.0}}
+            else:
+                rouge = Rouge()
+                scores = rouge.get_scores(" ".join(hypothesis), " ".join(reference))
+                result = scores[0]
+
+            for k, v in result.items():
+                score_dict[k].append(round(v["f"] * 100, 4))
+
+            bleu_score = sentence_bleu([list(label)], list(pred), smoothing_function=SmoothingFunction().method3)
+            score_dict["bleu-4"].append(round(bleu_score * 100, 4))
+
+        return {k: float(np.mean(v)) for k, v in score_dict.items()}

src/llmtuner/train/sftmm/trainer.py

@@ -0,0 +1,137 @@
+import json
+import os
+from types import MethodType
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from transformers import Seq2SeqTrainer
+
+from ...extras.constants import IGNORE_INDEX
+from ...extras.logging import get_logger
+from ..utils import create_custom_optimzer, create_custom_scheduler
+
+if TYPE_CHECKING:
+    from transformers.trainer import PredictionOutput
+    from peft import PeftModelForCausalLM
+    from ...hparams import FinetuningArguments
+
+logger = get_logger(__name__)
+
+
+class CustomSeq2SeqTrainer(Seq2SeqTrainer):
+    r"""
+    Inherits Seq2SeqTrainer to compute generative metrics such as BLEU and ROUGE.
+    """
+
+    def __init__(self, finetuning_args: "FinetuningArguments", **kwargs) -> None:
+        super().__init__(**kwargs)
+        self.finetuning_args = finetuning_args
+        if finetuning_args.use_badam:
+            from badam import clip_grad_norm_for_sparse_tensor
+
+            self.accelerator.clip_grad_norm_ = MethodType(clip_grad_norm_for_sparse_tensor, self.accelerator)
+
+    # def compute_loss(self, model, inputs, return_outputs=False):
+    #     print(inputs.keys())
+    #     device = "cuda"
+    #     input_ids = inputs.get("input_ids").to(device)
+    #     pixel_values = inputs.get("pixel_values").to(device, torch.float16)
+    #     attention_mask = inputs.get("attention_mask").to(device)
+    #     labels = inputs.get("labels").to(device)
+    #
+    #     outputs = model(input_ids=input_ids,
+    #                     pixel_values=pixel_values,
+    #                     labels=labels,
+    #                     # attention_mask=attention_mask,
+    #                     )
+    #     loss = outputs.loss
+    #     print("Loss:", loss.item())
+    #     return (loss, outputs) if return_outputs else loss
+
+    def create_optimizer(self) -> "torch.optim.Optimizer":
+        if self.optimizer is None:
+            self.optimizer = create_custom_optimzer(self.model, self.args, self.finetuning_args)
+        return super().create_optimizer()
+
+    def create_scheduler(
+            self, num_training_steps: int, optimizer: Optional["torch.optim.Optimizer"] = None
+    ) -> "torch.optim.lr_scheduler.LRScheduler":
+        create_custom_scheduler(self.args, num_training_steps, optimizer)
+        return super().create_scheduler(num_training_steps, optimizer)
+
+    def prediction_step(
+            self,
+            model: "torch.nn.Module",
+            inputs: Dict[str, Union[torch.Tensor, Any]],
+            prediction_loss_only: bool,
+            ignore_keys: Optional[List[str]] = None,
+    ) -> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
+        r"""
+        Removes the prompt part in the generated tokens.
+
+        Subclass and override to inject custom behavior.
+        """
+        labels = inputs["labels"].detach().clone() if "labels" in inputs else None  # backup labels
+        if self.args.predict_with_generate:
+            assert self.tokenizer.padding_side == "left", "This method only accepts left-padded tensor."
+            prompt_len, label_len = inputs["input_ids"].size(-1), inputs["labels"].size(-1)
+            if prompt_len > label_len:
+                inputs["labels"] = self._pad_tensors_to_target_len(inputs["labels"], inputs["input_ids"])
+            if label_len > prompt_len:  # truncate the labels instead of padding the inputs (llama2 fp16 compatibility)
+                inputs["labels"] = inputs["labels"][:, :prompt_len]
+
+        loss, generated_tokens, _ = super().prediction_step(  # ignore the returned labels (may be truncated)
+            model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys
+        )
+        if generated_tokens is not None and self.args.predict_with_generate:
+            generated_tokens[:, :prompt_len] = self.tokenizer.pad_token_id
+            generated_tokens = generated_tokens.contiguous()
+
+        return loss, generated_tokens, labels
+
+    def _pad_tensors_to_target_len(self, src_tensor: torch.Tensor, tgt_tensor: torch.Tensor) -> torch.Tensor:
+        r"""
+        Pads the tensor to the same length as the target tensor.
+        """
+        assert self.tokenizer.pad_token_id is not None, "Pad token is required."
+        padded_tensor = self.tokenizer.pad_token_id * torch.ones_like(tgt_tensor)
+        padded_tensor[:, -src_tensor.shape[-1]:] = src_tensor  # adopt left-padding
+        return padded_tensor.contiguous()  # in contiguous memory
+
+    def save_predictions(self, predict_results: "PredictionOutput") -> None:
+        r"""
+        Saves model predictions to `output_dir`.
+
+        A custom behavior that not contained in Seq2SeqTrainer.
+        """
+        if not self.is_world_process_zero():
+            return
+
+        output_prediction_file = os.path.join(self.args.output_dir, "generated_predictions.jsonl")
+        logger.info(f"Saving prediction results to {output_prediction_file}")
+
+        labels = np.where(
+            predict_results.label_ids != IGNORE_INDEX, predict_results.label_ids, self.tokenizer.pad_token_id
+        )
+        preds = np.where(
+            predict_results.predictions != IGNORE_INDEX, predict_results.predictions, self.tokenizer.pad_token_id
+        )
+
+        for i in range(len(preds)):
+            pad_len = np.nonzero(preds[i] != self.tokenizer.pad_token_id)[0]
+            if len(pad_len):
+                preds[i] = np.concatenate(
+                    (preds[i][pad_len[0]:], preds[i][: pad_len[0]]), axis=-1
+                )  # move pad token to last
+
+        decoded_labels = self.tokenizer.batch_decode(
+            labels, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )
+        decoded_preds = self.tokenizer.batch_decode(preds, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+
+        with open(output_prediction_file, "w", encoding="utf-8") as writer:
+            res: List[str] = []
+            for label, pred in zip(decoded_labels, decoded_preds):
+                res.append(json.dumps({"label": label, "predict": pred}, ensure_ascii=False))
+            writer.write("\n".join(res))

src/llmtuner/train/sftmm/workflow.py

@@ -0,0 +1,105 @@
+# 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 ..utils import create_modelcard_and_push
+from .metric import ComputeMetrics
+from .trainer import CustomSeq2SeqTrainer
+from .collator import DataCollatorForVis2Seq, ImageCaptioningDataset
+
+if TYPE_CHECKING:
+    from transformers import Seq2SeqTrainingArguments, TrainerCallback
+
+    from ...hparams import DataArguments, FinetuningArguments, GeneratingArguments, ModelArguments
+
+
+def run_sft_mm(
+        model_args: "ModelArguments",
+        data_args: "DataArguments",
+        training_args: "Seq2SeqTrainingArguments",
+        finetuning_args: "FinetuningArguments",
+        generating_args: "GeneratingArguments",
+        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)
+    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)
+    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
+
+    # Initialize our Trainer
+    trainer = CustomSeq2SeqTrainer(
+        model=model,
+        args=training_args,
+        finetuning_args=finetuning_args,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+        callbacks=callbacks,
+        compute_metrics=ComputeMetrics(tokenizer) if training_args.predict_with_generate else None,
+        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
+    gen_kwargs["pad_token_id"] = tokenizer.pad_token_id
+    gen_kwargs["logits_processor"] = get_logits_processor()
+
+    # Training
+    if training_args.do_train:
+        train_result = trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
+        trainer.save_model()
+        trainer.log_metrics("train", train_result.metrics)
+        trainer.save_metrics("train", train_result.metrics)
+        trainer.save_state()
+        if trainer.is_world_process_zero() and finetuning_args.plot_loss:
+            plot_loss(training_args.output_dir, keys=["loss", "eval_loss"])
+
+    # Evaluation
+    if training_args.do_eval:
+        metrics = trainer.evaluate(metric_key_prefix="eval", **gen_kwargs)
+        if training_args.predict_with_generate:  # eval_loss will be wrong if predict_with_generate is enabled
+            metrics.pop("eval_loss", None)
+        trainer.log_metrics("eval", metrics)
+        trainer.save_metrics("eval", metrics)
+
+    # Predict
+    if training_args.do_predict:
+        predict_results = trainer.predict(dataset, metric_key_prefix="predict", **gen_kwargs)
+        if training_args.predict_with_generate:  # predict_loss will be wrong if predict_with_generate is enabled
+            predict_results.metrics.pop("predict_loss", None)
+        trainer.log_metrics("predict", predict_results.metrics)
+        trainer.save_metrics("predict", predict_results.metrics)
+        trainer.save_predictions(predict_results)
+
+    # Create model card
+    create_modelcard_and_push(trainer, model_args, data_args, training_args, finetuning_args)

src/llmtuner/train/tuner.py

@@ -14,12 +14,11 @@ from .ppo import run_ppo
 from .pt import run_pt
 from .rm import run_rm
 from .sft import run_sft
-
+from .sftmm import run_sft_mm
 
 if TYPE_CHECKING:
     from transformers import TrainerCallback
 
-
 logger = get_logger(__name__)
 
 
@@ -31,6 +30,8 @@ def run_exp(args: Optional[Dict[str, Any]] = None, callbacks: Optional[List["Tra
         run_pt(model_args, data_args, training_args, finetuning_args, callbacks)
     elif finetuning_args.stage == "sft":
         run_sft(model_args, data_args, training_args, finetuning_args, generating_args, callbacks)
+    elif finetuning_args.stage == "sft_mm":
+        run_sft_mm(model_args, data_args, training_args, finetuning_args, generating_args, callbacks)
     elif finetuning_args.stage == "rm":
         run_rm(model_args, data_args, training_args, finetuning_args, callbacks)
     elif finetuning_args.stage == "ppo":