merge data part to the text stream

Former-commit-id: 7ee20286d9bcc2d5378bfd6bb02cd3648396d873

src/llmtuner/train/sftmm/collator.py

@@ -19,7 +19,9 @@ class DataCollatorForVis2Seq:
             texts.append(text)
             images.append(example["images"][0])
 
-        batch = self.processor(text=texts, images=images, return_tensors="pt", padding=True)
+        batch = self.processor(
+            text=texts, images=images, return_tensors="pt", padding=True
+        )
 
         labels = batch["input_ids"].clone()
         if self.processor.tokenizer.pad_token_id is not None:
@@ -27,3 +29,14 @@ class DataCollatorForVis2Seq:
         batch["labels"] = labels
 
         return batch
+
+
+@dataclass
+class DataCollatorForMLLM:
+    processor: AutoProcessor
+
+    def __call__(self, examples):
+        print(examples[0].keys())
+        print(examples[0]["input_ids"])
+        batch = {}
+        return batch

src/llmtuner/train/sftmm/workflow.py

@@ -1,47 +1,66 @@
 # 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
-from ...data import split_dataset, get_mm_dataset
+from ...data import get_dataset
 from ...extras.misc import get_logits_processor
 from ...extras.ploting import plot_loss
-from ...model import load_tokenizer, load_processor, load_model
+from ...model import load_processor, load_model
 from ..utils import create_modelcard_and_push
 from .metric import ComputeMetrics
 from .trainer import CustomSeq2SeqTrainer
-from .collator import DataCollatorForVis2Seq
+from transformers import DataCollatorForSeq2Seq
+from ...extras.constants import IGNORE_INDEX
 
 if TYPE_CHECKING:
     from transformers import Seq2SeqTrainingArguments, TrainerCallback
 
-    from ...hparams import DataArguments, FinetuningArguments, GeneratingArguments, ModelArguments
+    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,
+    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 = 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
-    model = load_model(processor.tokenizer, model_args, finetuning_args, training_args.do_train)
-    dataset = get_mm_dataset(processor, model_args, data_args, training_args, stage="sft")
+    tokenizer = processor.tokenizer
+    dataset = get_dataset(
+        tokenizer, model_args, data_args, training_args, "sft", processor
+    )
+    model = load_model(tokenizer, model_args, finetuning_args, training_args.do_train)
     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
+        setattr(
+            model, "_hf_peft_config_loaded", True
+        )  # hack here: make model compatible with prediction
     train_dataset = dataset
     eval_dataset = dataset
-    data_collator = DataCollatorForVis2Seq(
-        processor=processor,
+    data_collator = DataCollatorForSeq2Seq(
+        tokenizer=tokenizer,
+        pad_to_multiple_of=(
+            8 if tokenizer.padding_side == "right" else None
+        ),  # for shift short attention
+        label_pad_token_id=(
+            IGNORE_INDEX
+            if data_args.ignore_pad_token_for_loss
+            else tokenizer.pad_token_id
+        ),
     )
 
     # 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.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
@@ -52,19 +71,26 @@ def run_sft_mm(
         tokenizer=tokenizer,
         data_collator=data_collator,
         callbacks=callbacks,
-        compute_metrics=ComputeMetrics(tokenizer) if training_args.predict_with_generate else None,
+        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["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)
+        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)
@@ -75,19 +101,27 @@ def run_sft_mm(
     # 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
+        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 = 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)
+    create_modelcard_and_push(
+        trainer, model_args, data_args, training_args, finetuning_args
+    )