| from typing import Dict, Optional, List, Literal |
| from dataclasses import dataclass, field |
| import json |
| import transformers |
|
|
| from utils.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) |
| cut_embed_dim: Optional[int] = field( |
| default=3072, |
| metadata={"help": "cut off embed dim for tbstars model"}, |
| ) |
| parscale_n: Optional[int] = field( |
| default=2, |
| metadata={"help": "Use knn for partial fc"}, |
| ) |
| parscale_n_tokens: Optional[int] = field( |
| default=42, |
| metadata={"help": "parscale_n_tokens"}, |
| ) |
| parscale_attn_smooth: Optional[float] = field( |
| default=0.01, |
| metadata={"help": "parscale_attn_smooth"}, |
| ) |
| parscale_layer: Optional[int] = field( |
| default=1, |
| metadata={"help": "when to begin the parscale"}, |
| ) |
| perceiver_num_heads: Optional[int] = field( |
| default=32, |
| metadata={"help": "qformer_head"}, |
| ) |
| perceiver_num_queries: Optional[int] = field( |
| default=144, |
| metadata={"help": "qformer_dim"}, |
| ) |
| lora_load: Optional[bool] = field( |
| default=False, |
| metadata={"help": "if loading the lora pth"}, |
| ) |
| lora_pth: Optional[str] = field( |
| default='', |
| metadata={"help": "lora pth"}, |
| ) |
|
|
|
|
|
|
| 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."} |
| ) |
| tables: str = field( |
| default=None, metadata={"help": "Odps Tables."} |
| ) |
| tables_eval: str = field( |
| default=None, metadata={"help": "Odps Tables for eval."} |
| ) |
| max_lines: int = field( |
| default=-1, metadata={"help": "Max data lines."} |
| ) |
| max_lines_eval: int = field( |
| default=-1, metadata={"help": "Max data lines of eval."} |
| ) |
| image_type: str = field( |
| default="crop", metadata={"help": "eval instruction."} |
| ) |
| text_require: Optional[bool] = field( |
| default=False, |
| metadata={"help": "requrei text or not"}, |
| ) |
| |
|
|
|
|
| @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)." |
| }, |
| ) |
| eval_json: str = field( |
| default='eval_results.jsonl', metadata={"help": "json file to save eval result"} |
| ) |
|
|
| sub_loss: Optional[bool] = field( |
| default=True, |
| metadata={"help": "if use subspace loss"}, |
| ) |
| sub_loss_weight: Optional[float] = field( |
| default=1, |
| metadata={"help": "sub loss weight"}, |
| ) |
| margin_value: Optional[float] = field( |
| default=1, |
| metadata={"help": "margin value of triplet loss"}, |
| ) |
| triplet_weight: Optional[float] = field( |
| default=1, |
| metadata={"help": "margin value of triplet loss"}, |
| ) |
| club_lr: float = field( |
| default=1e-5, metadata={"help": "lambda for uniform loss"} |
| ) |
| club_weight_decay: float = field( |
| default=0.01, metadata={"help": "lambda for uniform loss"} |
| ) |
| club_flag: bool = field( |
| default=False, metadata={"help": "lambda for uniform loss"} |
| ) |
| club_iters: int = field( |
| default=5, metadata={"help": "lambda for uniform loss"} |
| ) |
| club_weight: float = field( |
| default=0.0001, metadata={"help": "lambda for uniform loss"} |
| ) |
| club_target: bool = field( |
| default=False, metadata={"help": "if regulation for target"} |
| ) |
| club_warmup_iters: int = field( |
| default=10, metadata={"help": "warm_up_iters for club"} |
| ) |
| use_flash_attn: bool = False |
| train_vision_encoder: bool = False |
| train_vision_projector: bool = False |
| train_llm: bool = False |
| train_parscale: bool = False |
| vision_projector_lr: float = None |
|
|
| def __post_init__(self): |
| super().__post_init__() |
| if self.resume_from_checkpoint == "False": |
| self.resume_from_checkpoint = False |
| elif self.resume_from_checkpoint == "True": |
| self.resume_from_checkpoint = True |
| 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 |
|
|
| @dataclass |
| class FCArguments: |
| r""" |
| Arguments pertaining to which techniques we are going to fine-tuning with. |
| """ |
| stage: Optional[Literal["classify_infer_crop", "classify_infer", "pt", "sft", "rm", "ppo", "retrieval", "reranker", "retrieval_infer_query", "retrieval_infer_item", "retrieval_hard_negative", "classify"]] = field( |
| default="classify", |
| metadata={"help": "Which stage will be performed in training."}, |
| ) |
| embed_dim: Optional[int] = field( |
| default=512, |
| metadata={"help": "embed dimension of output token"}, |
| ) |
| temperature: Optional[float] = field( |
| default=0.02, |
| metadata={"help": "temperature of the infonce loss"}, |
| ) |
| neg_gather_size: Optional[int] = field( |
| default=1, |
| metadata={ |
| "help": "how many gpus share negatives in the infonce loss. " |
| "1 = only in-device batch; 8 = share within a node; " |
| "-1 or world_size = share across all gpus. " |
| "candidate pool = 2 * neg_gather_size * per_device_train_batch_size" |
| }, |
| ) |
| classify_num: Optional[int] = field( |
| default=100000, |
| metadata={"help": "Num of classification"}, |
| ) |
| classify_sample_rate: Optional[float] = field( |
| default=0.05, |
| metadata={"help": "Sample rate for partial fc"}, |
| ) |
| classify_use_knn: Optional[bool] = field( |
| default=True, |
| metadata={"help": "Use knn for partial fc"}, |
| ) |
| classify_global_knn: Optional[bool] = field( |
| default=False, |
| metadata={"help": "Use knn for partial fc"}, |
| ) |
| classify_knn_topk: Optional[int] = field( |
| default=30, |
| metadata={"help": "Knn topk for partial fc"}, |
| ) |
| classify_knn_chunk_size: Optional[int] = field( |
| default=5000, |
| metadata={"help": "knn chunk size for sim matrix calc: only open when global_knn"}, |
| ) |
| classify_knn_ratio: Optional[float] = field( |
| default=0.5, |
| metadata={"help": "Knn rate for partial fc"}, |
| ) |
| classify_knn_weight_cover_rate: Optional[float] = field( |
| default=0.3, |
| metadata={"help": "Knn rate for partial fc"}, |
| ) |
| classify_knn_update_after_save: Optional[bool] = field( |
| default=True, |
| metadata={"help": "knn_update_after_save"}, |
| ) |
| classify_loss_type: Optional[str] = field( |
| default="arcface", |
| metadata={"help": "loss function for partial fc"}, |
| ) |
| classify_ckpt_part: Optional[bool] = field( |
| default=True, |
| metadata={"help": "partial fc ckpt save load partly"}, |
| ) |
| classify_freeze_backbone: Optional[bool] = field( |
| default=False, |
| metadata={"help": "freeze backbone, only train partial fc"}, |
| ) |
| model_name_or_path_pfc: Optional[str] = field( |
| default="/data/oss_bucket_0/guijiang/outputs/tongkuan/lamra_qwen2_5_fc_debug", |
| metadata={"help": "pth to fc pth"}, |
| ) |
| model_name_or_path_proj: Optional[str] = field( |
| default="/data/oss_bucket_0/guijiang/outputs/tongkuan/lamra_qwen2_5_fc_debug", |
| metadata={"help": "pth to fc pth"}, |
| ) |
| proj_name: Optional[Literal["load_add_projector", "load_mlp_projector"]] = field( |
| default="load_add_projector", |
| metadata={"help": "Which stage will be performed in training."}, |
| ) |
| pfc_trainable: Optional[bool] = field( |
| default=True, |
| metadata={"help": "train pfc or not"}, |
| ) |
| def save_to_json(self, json_path: str): |
| r"""Saves the content of this instance in JSON format inside `json_path`.""" |
| json_string = json.dumps(asdict(self), indent=2, sort_keys=True) + "\n" |
| with open(json_path, "w", encoding="utf-8") as f: |
| f.write(json_string) |
|
|
| @classmethod |
| def load_from_json(cls, json_path: str): |
| r"""Creates an instance from the content of `json_path`.""" |
| with open(json_path, "r", encoding="utf-8") as f: |
| text = f.read() |
|
|
| return cls(**json.loads(text)) |
|
|
