 b/src/train/train.py
+import logging
+from dataclasses import dataclass, field
+from typing import Optional
+import transformers
+from peft import LoraConfig, get_peft_model
+from transformers import (
+    Trainer,
+    TrainingArguments,
+    AutoTokenizer,
+    PreTrainedTokenizer
+)
+from src.dataset.collator import InstructionTuningCollator
+from src.dataset.dataset import InstructionTuningDataset
+from src.model.init_llemr import init_llemr
+from src.model.modeling_llemr import LlemrForConditionalGeneration
+from src.model.utils import find_all_linear_names
+logger = logging.getLogger(__name__)
+@dataclass
+class ModelArguments:
+    name_or_path: Optional[str] = field(default=None)
+    llm_pretrained_model_name_or_path: Optional[str] = field(default="Qwen/Qwen2-0.5B-Instruct")
+    train_type: Optional[str] = field(
+        default="train_both",
+        metadata={
+            "help": """
+. train_multi_modal_projector
+. train_both
+            """
+        },
+    )
+    use_lora: Optional[bool] = field(default=True)
+    lora_r: int = 32
+    lora_alpha: int = 16
+    lora_dropout: float = 0.05
+    lora_bias: str = "none"
+    vision_hidden_size: int = 768
+@dataclass
+class DataArguments:
+    source: Optional[str] = field(default="note")
+def load_model(model_args: ModelArguments):
+    if model_args.name_or_path is not None:
+        logging.warning(f"Load model {model_args.name_or_path} from pretrained")
+        model = LlemrForConditionalGeneration.from_pretrained(
+            model_args.name_or_path
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_args.name_or_path,
+            padding_side="left"
+        )
+    else:
+        logging.warning(f"Init model {model_args.llm_pretrained_model_name_or_path}")
+        model, tokenizer = init_llemr(
+            model_args.llm_pretrained_model_name_or_path, model_args.vision_hidden_size
+        )
+    assert model_args.train_type in ["train_multi_modal_projector", "train_both"]
+    if model_args.train_type == "train_multi_modal_projector":
+        logging.warning("Train multi_modal_projector")
+        for param in model.language_model.parameters():
+            param.requires_grad = False
+    else:
+        logging.warning("Train both")
+    if model_args.use_lora:
+        assert model_args.train_type == "train_both"
+        logging.warning("Use Lora")
+        config = LoraConfig(
+            r=model_args.lora_r,
+            lora_alpha=model_args.lora_alpha,
+            target_modules=find_all_linear_names(model),
+            lora_dropout=model_args.lora_dropout,
+            bias=model_args.lora_bias,
+            task_type="CAUSAL_LM",
+            modules_to_save=["multi_modal_projector"],
+        )
+        model = get_peft_model(model, config)
+    else:
+        logging.warning("Not use Lora")
+    return model, tokenizer
+def load_data(data_args: DataArguments, tokenizer: PreTrainedTokenizer):
+    train_dataset = InstructionTuningDataset(
+        split="train",
+        source=data_args.source,
+    )
+    val_dataset = InstructionTuningDataset(
+        split="val",
+        source=data_args.source,
+    )
+    collator = InstructionTuningCollator(
+        tokenizer=tokenizer,
+    )
+    return train_dataset, val_dataset, collator
+def train():
+    parser = transformers.HfArgumentParser(
+        (ModelArguments, DataArguments, TrainingArguments)
+    )
+    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+    model, tokenizer = load_model(model_args)
+    train_dataset, val_dataset, collator = load_data(data_args, tokenizer)
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=val_dataset,
+        data_collator=collator,
+    )
+    tokenizer.save_pretrained(training_args.output_dir)
+    trainer.train()
+if __name__ == "__main__":
+    train()