| from typing import Dict, Optional, List |
| from dataclasses import dataclass, field |
|
|
| import transformers |
|
|
| from supported_models import MODEL_HF_PATH, MODEL_FAMILIES |
|
|
|
|
| @dataclass |
| class ModelArguments: |
| model_id: str = field(default="llava-1.5-7b") |
| model_local_path: Optional[str] = field(default=None) |
|
|
| def __post_init__(self): |
| assert self.model_id in MODEL_HF_PATH, f"Unknown model_id: {self.model_id}" |
| self.model_hf_path: str = MODEL_HF_PATH[self.model_id] |
| assert self.model_id in MODEL_FAMILIES, f"Unknown model_id: {self.model_id}" |
| self.model_family_id: str = MODEL_FAMILIES[self.model_id] |
|
|
| if not self.model_local_path: |
| self.model_local_path = self.model_hf_path |
|
|
|
|
| @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."} |
| ) |
| 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") |
| image_data_path: str = field( |
| default=None, metadata={"help": "Path to the image data json file."} |
| ) |
| text_data_path: str = field( |
| default=None, metadata={"help": "Path to the text data json file."} |
| ) |
| query_data_path: str = field( |
| default=None, metadata={"help": "Path to the query data json file."} |
| ) |
| cand_pool_path: str = field( |
| default=None, metadata={"help": "Path to the cand pool data json file."} |
| ) |
| instructions_path: str = field( |
| default=None, metadata={"help": "Path to the instructions data json file."} |
| ) |
| rerank_data_path: str = field( |
| default=None, metadata={"help": "Path to the rerank data json file."} |
| ) |
| image_path_prefix: str = field( |
| default=None, metadata={"help": "Path to the image files."} |
| ) |
| train_data_path: str = field( |
| default=None, metadata={"help": "Path to the rerank data json file."} |
| ) |
| eval_data_path: str = field( |
| default=None, metadata={"help": "Path to the rerank data json file."} |
| ) |
| num_train_docs: int = field(default=10) |
| full_crop_mode: bool = False |
| class_mode: int = field(default=0) |
| use_crop: bool = False |
| dataset_type: str = "TongkuanPointWiseDataset" |
| format_prompt: str = "format_prompt/short_cot_format.jinja" |
| max_full_tokens: int = field(default=20) |
| max_crop_tokens: int = field(default=60) |
| max_visual_tokens: int = field(default=100) |
|
|
|
|
| @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 = False |
| train_vision_encoder: bool = False |
| train_vision_projector: bool = False |
| vision_projector_lr: float = None |
| tp_size: int = field(default=1) |
|
|
| def __post_init__(self): |
| super().__post_init__() |
| self.remove_unused_columns = False |
|
|
|
|
| @dataclass |
| class LoraArguments: |
| use_lora: bool = True |
| use_vision_lora: bool = True |
| q_lora: bool = False |
| lora_r: int = 64 |
| lora_alpha: int = 16 |
| lora_dropout: float = 0.05 |
| lora_weight_path: str = "" |
| lora_bias: str = "none" |
| vision_lora_r: int = 16 |
| vision_lora_alpha: int = 16 |
|
|
