File size: 3,144 Bytes
0ab2566 | 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 | from dataclasses import dataclass, field
from typing import Optional
@dataclass
class CoTArguments:
r"""
CoT embedding 训练特有的参数，与 utils.arguments 中已有的参数类配合使用。
"""
infonce_weight: Optional[float] = field(
default=1.0,
metadata={"help": "weight of the embedding infonce loss"},
)
cot_weight: Optional[float] = field(
default=1.0,
metadata={"help": "weight of the cot cross entropy loss"},
)
max_cot_length: Optional[int] = field(
default=192,
metadata={"help": "max token length of the cot text, longer cot will be truncated"},
)
cot_emb_suffix: Optional[str] = field(
default="<|im_end|>",
metadata={"help": "suffix appended after the cot text, its hidden state is pooled as query embedding"},
)
doc_append_eos: Optional[bool] = field(
default=False,
metadata={"help": "also append cot_emb_suffix to doc side so that query/doc pool on the same token"},
)
train_lm_head: Optional[bool] = field(
default=False,
metadata={
"help": "whether to finetune lm_head. lm_head is weight-tied with embed_tokens: "
"under full finetuning both are unfrozen/frozen together; "
"under LoRA a decoupled copy of lm_head is trained via modules_to_save "
"while embed_tokens stays frozen"
},
)
query_image_type: Optional[str] = field(
default=None,
metadata={"help": "image type of query side: full / crop / both. fall back to data_args.image_type"},
)
doc_image_type: Optional[str] = field(
default=None,
metadata={"help": "image type of doc side: full / crop / both. fall back to data_args.image_type"},
)
# ------------------------------------------------------------------
# wandb 设置：登录 key / project 等在此填写，use_wandb 打开后生效
# ------------------------------------------------------------------
use_wandb: Optional[bool] = field(
default=False,
metadata={"help": "enable wandb logging"},
)
wandb_api_key: Optional[str] = field(
default="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", # TODO: 在此填写 wandb 登录 key
metadata={"help": "wandb login api key"},
)
wandb_project: Optional[str] = field(
default="cot_embedding", # TODO: 在此填写 wandb project 名称
metadata={"help": "wandb project name"},
)
wandb_entity: Optional[str] = field(
default="", # 可选：wandb 团队/用户名，不填用账号默认值
metadata={"help": "wandb entity (team or username), optional"},
)
wandb_run_name: Optional[str] = field(
default="lora_alpha128_train_head_local",
metadata={"help": "wandb run name, fall back to training_args.run_name"},
)
wandb_mode: Optional[str] = field(
default="online",
metadata={"help": "wandb mode: online / offline, use offline when workers have no internet access"},
)
|