code: def train():
      global local_rank
  
      parser = transformers.HfArgumentParser(
          (ModelArguments, DataArguments, TrainingArguments)
      )
      (
          model_args,
          data_args,
          training_args,
      ) = parser.parse_args_into_dataclasses()
      local_rank = training_args.local_rank
      compute_dtype = (
          torch.float16
          if training_args.fp16
          else (torch.bfloat16 if training_args.bf16 else torch.float32)
      )
  
      bnb_model_from_pretrained_args = {}
      if training_args.bits in [4, 8]:
          from transformers import BitsAndBytesConfig
  
          bnb_model_from_pretrained_args.update(
              dict(
                  device_map={"": training_args.device},
                  load_in_4bit=training_args.bits == 4,
                  load_in_8bit=training_args.bits == 8,
                  quantization_config=BitsAndBytesConfig(
                      load_in_4bit=training_args.bits == 4,
                      load_in_8bit=training_args.bits == 8,
                      llm_int8_threshold=6.0,
                      llm_int8_has_fp16_weight=False,
                      bnb_4bit_compute_dtype=compute_dtype,
                      bnb_4bit_use_double_quant=training_args.double_quant,
                      bnb_4bit_quant_type=training_args.quant_type,  # {'fp4', 'nf4'}
                  ),
              )
          )
  
      if model_args.vision_tower is not None:
          if "mpt" in model_args.model_name_or_path:
              config = transformers.AutoConfig.from_pretrained(
                  model_args.model_name_or_path, trust_remote_code=True
              )
              config.attn_config[
                  "attn_impl"
              ] = training_args.mpt_attn_impl
              model = LlavaMPTForCausalLM.from_pretrained(
                  model_args.model_name_or_path,
                  config=config,
                  cache_dir=training_args.cache_dir,
                  **bnb_model_from_pretrained_args,
              )
          else:
              model = LlavaLlamaForCausalLM.from_pretrained(
                  model_args.model_name_or_path,
                  cache_dir=training_args.cache_dir,
                  **bnb_model_from_pretrained_args,
              )
      else:
          model = transformers.LlamaForCausalLM.from_pretrained(
              model_args.model_name_or_path,
              cache_dir=training_args.cache_dir,
              **bnb_model_from_pretrained_args,
          )
      model.config.use_cache = False
  
      if model_args.freeze_backbone:
          model.model.requires_grad_(False)
  
      if training_args.bits in [4, 8]:
          from peft import prepare_model_for_kbit_training
  
          model.config.torch_dtype = (
              torch.float32
              if training_args.fp16
              else (
                  torch.bfloat16
                  if training_args.bf16
                  else torch.float32
              )
          )
          model = prepare_model_for_kbit_training(
              model,
              use_gradient_checkpointing=training_args.gradient_checkpointing,
          )
  
      if training_args.gradient_checkpointing:
          if hasattr(model, "enable_input_require_grads"):
              model.enable_input_require_grads()
          else:
  
              def make_inputs_require_grad(module, input, output):
                  output.requires_grad_(True)
  
              model.get_input_embeddings().register_forward_hook(
                  make_inputs_require_grad
              )
  
      if training_args.lora_enable:
          from peft import LoraConfig, get_peft_model
  
          lora_config = LoraConfig(
              r=training_args.lora_r,
              lora_alpha=training_args.lora_alpha,
              target_modules=find_all_linear_names(model),
              lora_dropout=training_args.lora_dropout,
              bias=training_args.lora_bias,
              task_type="CAUSAL_LM",
          )
          if training_args.bits == 16:
              if training_args.bf16:
                  model.to(torch.bfloat16)
              if training_args.fp16:
                  model.to(torch.float16)
          rank0_print("Adding LoRA adapters...")
          model = get_peft_model(model, lora_config)
  
      if "mpt" in model_args.model_name_or_path:
          tokenizer = transformers.AutoTokenizer.from_pretrained(
              model_args.model_name_or_path,
              cache_dir=training_args.cache_dir,
              model_max_length=training_args.model_max_length,
              padding_side="right",
          )
      else:
          tokenizer = transformers.AutoTokenizer.from_pretrained(
              model_args.model_name_or_path,
              cache_dir=training_args.cache_dir,
              model_max_length=training_args.model_max_length,
              padding_side="right",
              use_fast=False,
          )
  
      if model_args.version == "v0":
          if tokenizer.pad_token is None:
              smart_tokenizer_and_embedding_resize(
                  special_tokens_dict=dict(pad_token="[PAD]"),
                  tokenizer=tokenizer,
                  model=model,
              )
      elif model_args.version == "v0.5":
          tokenizer.pad_token = tokenizer.unk_token
      else:
          tokenizer.pad_token = tokenizer.unk_token
          if model_args.version in conversation_lib.conv_templates:
              conversation_lib.default_conversation = (
                  conversation_lib.conv_templates[model_args.version]
              )
          else:
              conversation_lib.default_conversation = (
                  conversation_lib.conv_templates["vicuna_v1"]
              )
  
      if model_args.vision_tower is not None:
          model.get_model().initialize_vision_modules(
              model_args=model_args, fsdp=training_args.fsdp
          )
  
          vision_tower = model.get_vision_tower()
          vision_tower.to(
              dtype=torch.float16, device=training_args.device
          )
  
          data_args.image_processor = vision_tower.image_processor
          data_args.is_multimodal = True
  
          model.config.image_aspect_ratio = data_args.image_aspect_ratio
          model.config.image_grid_pinpoints = (
              data_args.image_grid_pinpoints
          )
  
          model.config.tune_mm_mlp_adapter = (
              training_args.tune_mm_mlp_adapter
          ) = model_args.tune_mm_mlp_adapter
          if model_args.tune_mm_mlp_adapter:
              model.requires_grad_(False)
              for p in model.get_model().mm_projector.parameters():
                  p.requires_grad = True
  
          model.config.freeze_mm_mlp_adapter = (
              training_args.freeze_mm_mlp_adapter
          )
          if training_args.freeze_mm_mlp_adapter:
              for p in model.get_model().mm_projector.parameters():
                  p.requires_grad = False
  
          if training_args.bits in [4, 8]:
              model.get_model().mm_projector.to(
                  dtype=compute_dtype, device=training_args.device
              )
  
          model.config.mm_use_output_start_end = (
              data_args.mm_use_output_start_end
          ) = model_args.mm_use_output_start_end
          training_args.use_output_start_end = (
              model_args.mm_use_output_start_end
          )
          model.config.mm_use_im_start_end = (
              data_args.mm_use_im_start_end
          ) = model_args.mm_use_im_start_end
          training_args.use_im_start_end = (
              model_args.mm_use_im_start_end
          )
          model.config.mm_use_im_patch_token = (
              model_args.mm_use_im_patch_token
          )
          model.initialize_vision_tokenizer(
              model_args, tokenizer=tokenizer
          )
  
      if training_args.bits in [4, 8]:
          from peft.tuners.lora import LoraLayer
  
          for name, module in model.named_modules():
              if isinstance(module, LoraLayer):
                  if training_args.bf16:
                      module = module.to(torch.bfloat16)
              if "norm" in name:
                  module = module.to(torch.float32)
              if "lm_head" in name or "embed_tokens" in name:
                  if hasattr(module, "weight"):
                      if (
                          training_args.bf16
                          and module.weight.dtype == torch.float32
                      ):
                          module = module.to(torch.bfloat16)
  
      data_module = make_supervised_data_module(
          tokenizer=tokenizer, data_args=data_args
      )
      trainer = LLaVATrainer(
          model=model,
          tokenizer=tokenizer,
          args=training_args,
          **data_module,
      )
  
      if list(
          pathlib.Path(training_args.output_dir).glob("checkpoint-*")
      ):
          trainer.train(resume_from_checkpoint=True)
      else:
          trainer.train()
      trainer.save_state()
  
      model.config.use_cache = True
  
      if training_args.lora_enable:
          state_dict = get_peft_state_maybe_zero_3(
              model.named_parameters(), training_args.lora_bias
          )
          non_lora_state_dict = get_peft_state_non_lora_maybe_zero_3(
              model.named_parameters()
          )
          if (
              training_args.local_rank == 0
              or training_args.local_rank == -1
          ):
              model.config.save_pretrained(training_args.output_dir)
              model.save_pretrained(
                  training_args.output_dir, state_dict=state_dict
              )
              torch.save(
                  non_lora_state_dict,
                  os.path.join(
                      training_args.output_dir,
                      "non_lora_trainables.bin",
                  ),
              )
      else:
          safe_save_model_for_hf_trainer(
              trainer=trainer, output_dir=training_args.output_dir
          )
  
imports: import os
  import copy
  from dataclasses import dataclass, field
  import json
  import logging
  import pathlib
  from typing import Dict, Optional, Sequence, List
  import torch
  import transformers
  from llmga.llava.constants import IGNORE_INDEX, IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN, DEFAULT_OUTPUT_START_TOKEN, DEFAULT_OUTPUT_END_TOKEN
  from torch.utils.data import Dataset
  from llmga.llava.train.llava_trainer import LLaVATrainer
  from llmga.llava import conversation
  from llmga.llava.model import *
  from llmga.llava.mm_utils import tokenizer_image_token
  from PIL import Image
  from llmga.llava.masks.make_mask import get_mask_generator
  from llmga.llava.prompt_temp import outpaint_prompt, inpaint_prompt, textextend_prompt, regen_prompt, ques_prompt, textextend_prompt2
  import random
  import numpy
  import lmdb
  import io
  import pytextrank
  from torchvision import transforms
label_desc: Use image node_id 19 for calling a global function with example usage: train() without return types