"""CoT embedding 本地评测参数文件。

在这里直接修改参数，然后运行：
    python cot_eval/eval_local.py
无需在命令行传任何参数。

注意：bool 参数直接写 True / False（不是字符串）。
"""
from dataclasses import dataclass
from typing import Optional


@dataclass
class EvalConfig:
    # ------------------------------------------------------------------
    # 模型
    # ------------------------------------------------------------------
    # base 模型路径（与训练时的 model_local_path 一致）
    model_id: str = "/data/oss_bucket_0/liqi/models/Qwen3-VL-Embedding-2B/"
    emb_dim: int = 2048
    out_dim: int = 2048

    # 训练产出的 checkpoint：
    # lora=True  -> lora_pth 为 LoRA adapter 目录（checkpoint-xxx），base 从 model_id 加载后 merge；
    # lora=False -> lora_pth 为全参微调 checkpoint 目录（checkpoint-xxx），直接加载完整权重；
    #               置空字符串则只评测 model_id 的 base 模型。
    lora: bool = False
    # all_eval=False 时给到单个 checkpoint-xxx 目录；
    # all_eval=True 时给到 run 根目录，会遍历其中所有 checkpoint-*
    ckpt: int = 2500
    root_pth: str = "/data/oss_bucket_0/liqi/cot_embedding/nebula_train/dapei_loraFalse_lr1e-5_btz256_ep4_candidate8_cot0.1_8gpu/train_qwen3vl_cot_lora_False_lmhead_True_epoch_4_lr1e-5_candidate8_cot0.1/"
    lora_pth: str = f"{root_pth}/checkpoint-{ckpt}"
    all_eval: bool = False

    # ------------------------------------------------------------------
    # 数据集（与 baseline 评测一致）
    # ------------------------------------------------------------------
    test_data_type: str = "dapei"
    test_json: str = f"/data/oss_bucket_0/liqi/dataset/embedding_pf/{test_data_type}_test.json"
    max_lines: int = 100000000
    image_type: str = "both"          # full / crop / both
    query_image_type: Optional[str] = None   # 不填则用 image_type
    doc_image_type: Optional[str] = None     # 不填则用 image_type
    # 本地测试图片目录：命中则本地读取，找不到时回退 OSS；置 None 则全走 OSS
    image_path_prefix: Optional[str] = "/data/oss_bucket_0/peifeng/multi_vlm2vec/dataset/merge/test_image"
    query_text_require: bool = True
    doc_text_require: bool = True
    batch_size: int = 256

    # ------------------------------------------------------------------
    # CoT 设置（需与训练配置保持一致）
    # ------------------------------------------------------------------
    prompt_index: int = 0
    max_cot_new_tokens: int = 192     # 与训练 max_cot_length 对齐
    cot_emb_suffix: str = "<|im_end|>"
    doc_append_eos: bool = False
    # query = prompt + cot，序列比纯 embedding 评测长
    model_max_length: int = 2048

    # ------------------------------------------------------------------
    # 结果保存
    # ------------------------------------------------------------------
    # 生成的 CoT 文本（qid -> cot），置空字符串则不保存
    # all_eval=True 时文件名自动加 checkpoint-xxx_ 前缀，每个 ckpt 一个文件
    save_cot_json: str = f"{root_pth}/eval_result/ckpt-{ckpt}/cot.json"
    # 指标结果，置空字符串则不保存
    # all_eval=True 时所有 ckpt 汇总进这一个文件：{checkpoint-xxx: results}
    result_json: str = f"{root_pth}/eval_result/ckpt-{ckpt}/metrics.json"
    train_data_type: str = ""