VersaViT/arguments.py at main · Tencent/VersaViT · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
from typing import Dict, Optional, List, Union
from dataclasses import dataclass, field

import transformers


@dataclass
class ModelArguments:
    model_id: str = field(default="llava-1.5-7b")
    model_local_path: Optional[str] = field(default=None)
    text_decoder_name: str = field(default='')
    text_encoder_name: str=field(default='')
    vit_path: str=field(default='')


@dataclass
class DataArguments:
    data_path: str = field(
        default=None, metadata={"help": "Path to the training data json file."}
    )
    eval_data_path: Optional[str] = field(
        default=None, metadata={"help": "Path to the evaluation data json file."}
    )

    train_data_path: Optional[str] = field(default=None, metadata={'help': "list of paths to the training data."})
    cap_train_data_path: Optional[str] = field(default=None, metadata={'help': "list of paths to the captioning training data."})
    seg_train_data_path: Optional[str] = field(default=None, metadata={'help': "list of paths to the segmentation training data."})
    depth_train_data_path: Optional[str] = field(default=None, metadata={'help': "list of paths to the depth training data."})

    sample_ratios: Optional[str] = field(default=None, metadata={'help': "sample_ratios for each dataset."})
    cap_sample_ratios: Optional[str] = field(default=None, metadata={'help': "sample_ratios for the captioning dataset."})
    seg_sample_ratios: Optional[str] = field(default=None, metadata={'help': "sample_ratios for the segmentation dataset."})
    depth_sample_ratios: Optional[str] = field(default=None, metadata={'help': "sample_ratios for the depth dataset."})

    image_folder: Optional[str] = field(default=None)
    video_folder: Optional[str] = field(default=None)
    num_frames: Optional[int] = field(default=8)
    user_key: Optional[str] = field(default="human")
    assistant_key: Optional[str] = field(default="gpt")

    captioning_files_per_tar: int = field(default=None)
    segmentation_files_per_tar: int = field(default=None)
    depth_files_per_tar: int = field(default=None)

    def __post_init__(self):
        self.train_data_path = self.train_data_path.split(',') if self.train_data_path else []
        self.cap_train_data_path = self.cap_train_data_path.split(',') if self.cap_train_data_path else []
        self.seg_train_data_path = self.seg_train_data_path.split(',') if self.seg_train_data_path else []
        self.depth_train_data_path = self.depth_train_data_path.split(',') if self.depth_train_data_path else []

        self.sample_ratios = [float(p) for p in self.sample_ratios.split(',')] if self.sample_ratios else []
        self.cap_sample_ratios = [float(p) for p in self.cap_sample_ratios.split(',')] if self.cap_sample_ratios else []
        self.seg_sample_ratios = [float(p) for p in self.seg_sample_ratios.split(',')] if self.seg_sample_ratios else []
        self.depth_sample_ratios = [float(p) for p in self.depth_sample_ratios.split(',')] if self.depth_sample_ratios else []


@dataclass
class TrainingArguments(transformers.TrainingArguments):
    model_max_length: int = field(
        default=1024,
        metadata={
            "help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
        },
    )
    use_flash_attn: bool = field(default=False)
    train_vision_encoder: bool = field(default=False)
    train_vision_projector: bool = field(default=False)
    mask_question_tokens: bool = field(default=True)
    num_samples: int = field(default=None)
    vision_lr: float = field(default=None)
    captioning_projection_lr: float = field(default=None)
    text_decoder_lr: float = field(default=None)
    freeze_text_decoder: bool = field(default=True)
    task_sequence: Optional[str] = field(default=None, metadata={'help': "list of paths to the training data."})
    captioning_num_samples: int = field(default=None)
    segmentation_num_samples: int = field(default=None)
    depth_num_samples: int = field(default=None)
    captioning_num_batches: int = field(default=None)
    segmentation_num_batches: int = field(default=None)
    depth_num_batches: int = field(default=None)
    captioning_head_pretrained_path: str = field(default=None)
    segmentation_batch_size: int = field(default=None)
    captioning_batch_size: int = field(default=None)
    depth_batch_size: int = field(default=None)
    captioning_loss_weight: float = field(default=1.0)
    segmentation_loss_weight: float = field(default=1.0)
    depth_loss_weight: float = field(default=1.0)


    def __post_init__(self):
        super().__post_init__()
        self.remove_unused_columns = False
        self.task_sequence = self.task_sequence.split(',') if self.task_sequence else []


@dataclass
class LoraArguments:
    use_lora: bool = field(default=True)
    use_vision_lora: bool = field(default=True)
    q_lora: bool = field(default=False)
    lora_r: int = field(default=8)
    lora_alpha: int = field(default=16)
    lora_dropout: float = field(default=0.05)
    lora_weight_path: str = ""
    lora_bias: str = "none"