| 53 | |
| 54 | @dataclass |
| 55 | class ModelConfig: |
| 56 | model_name_or_path: Optional[str] = None |
| 57 | model_revision: str = "main" |
| 58 | |
| 59 | output_dim: int = 1 |
| 60 | |
| 61 | use_special_tokens: bool = False |
| 62 | |
| 63 | freeze_vision_tower: bool = field(default=False) |
| 64 | freeze_llm: bool = field(default=False) |
| 65 | tune_merger: bool = field(default=False) |
| 66 | |
| 67 | torch_dtype: Optional[Literal["auto", "bfloat16", "float16", "float32"]] = None |
| 68 | trust_remote_code: bool = False |
| 69 | attn_implementation: Optional[str] = None |
| 70 | load_in_8bit: bool = False |
| 71 | load_in_4bit: bool = False |
| 72 | bnb_4bit_quant_type: Literal["fp4", "nf4"] = "nf4" |
| 73 | use_bnb_nested_quant: bool = False |
| 74 | reward_token: Literal["last", "mean", "special"] = "last" |
| 75 | loss_type: Literal["bt", "reg", "btt", "margin", "constant_margin", "scaled"] = "regular" |
| 76 | |
| 77 | def __post_init__(self): |
| 78 | if self.load_in_8bit and self.load_in_4bit: |
| 79 | raise ValueError("You can't use 8 bit and 4 bit precision at the same time") |
| 80 | |
| 81 | # if isinstance(self.lora_target_modules, list) and len(self.lora_target_modules) == 1: |
| 82 | # self.lora_target_modules = self.lora_target_modules[0] |
| 83 | |
| 84 | # if isinstance(self.lora_namespan_exclude, list) and len(self.lora_namespan_exclude) == 1: |
| 85 | # self.lora_namespan_exclude = self.lora_namespan_exclude[0] |
| 86 | |
| 87 | ########## Functions for get trainable modules' parameters ########## |
| 88 | |