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"},
    )