| 69 | |
| 70 | @dataclass |
| 71 | class ModelMergerConfig: |
| 72 | operation: str # 'merge' or 'test' |
| 73 | backend: str |
| 74 | local_dir: str |
| 75 | hf_model_config_path: str |
| 76 | target_dir: Optional[str] = "tmp" |
| 77 | hf_upload_path: Optional[str] = None |
| 78 | private: bool = False |
| 79 | test_hf_dir: Optional[str] = None |
| 80 | tie_word_embedding: bool = False |
| 81 | is_value_model: bool = False |
| 82 | hf_model_path: Optional[str] = None |
| 83 | hf_upload: bool = field(init=False) |
| 84 | |
| 85 | def __post_init__(self): |
| 86 | self.hf_upload = self.operation == "merge" and bool(self.hf_upload_path) |
| 87 | if self.operation == "test": |
| 88 | self.target_dir = None |
| 89 | self.hf_upload_path = None |
| 90 | self.private = False |
| 91 | |
| 92 | |
| 93 | class BaseModelMerger(ABC): |