# ================================================================
# PMPS 分析实验配置 (GPT-2, GSM8K)
# 用于 extract_latent_states.py --extract_pmps
# ================================================================
#
# 需要填写的参数只有 ckpt_dir，其余参数必须和训练时保持一致，
# 否则模型权重会加载失败或结构不匹配。
#
# 判断依据：如果训练时某参数改变了模型结构（如 num_latent, prj_dim,
# pmps_nmb_prototypes），那分析时必须用相同值。

# ==================== 必填：模型和 Checkpoint ====================
model_name_or_path: /data/oss_bucket_0/liqi/models/gpt2  # 基座模型路径
ckpt_dir: /data/oss_bucket_0/liqi/codi_pmps/k24_t0.1_mt0.1_fac0.5_l6_lr3e3_btz256/outputs/checkpoint-58000/              # 例: /data/.../checkpoint-15000

# ==================== 必须和训练一致的参数 ====================
# --- LoRA 结构 ---
lora_r: 128
lora_alpha: 32
lora_init: True
use_lora: True

# --- Latent 结构 ---
num_latent: 6              # 训练时用了几个 latent token
use_prj: True              # 是否用 projection layer
prj_dim: 768               # projection hidden dim（GPT-2 用 768）

# --- PMPS/PMPS 结构（决定了 pmps_module 的权重形状）---
pmps_loss_factor: 0.5      # > 0 表示有 pmps_module，分析时会提取 PMPS 内部状态
pmps_nmb_prototypes: 24    # prototype 数量 K
pmps_feat_dim: 128         # 投影特征维度 d
pmps_hidden_mlp: 512       # 投影 MLP 隐藏维度
pmps_use_projector: True   # 是否使用投影头

# --- PMPS 超参（不影响结构，但 forward 计算需要）---
pmps_temperature: 0.1
pmps_match_temperature: 0.3
pmps_epsilon: 0.05
pmps_sinkhorn_iterations: 3
pmps_layer: last

# --- PSA（不影响结构，分析时不使用）---
pmps_position_bias_max: 6.0
pmps_position_bias_schedule: cosine

# --- Top-K Embedding ---
enable_top_k: False
latent_top_k: 0
latent_top_k_temperature: 1.0

# --- 其他训练时的配置 ---
remove_eos: True
diversity_loss_factor: 0.0

# ==================== 分析时无关但 parser 需要的参数 ====================
data_name: gsm8k            # 决定加载哪个测试集
model_max_length: 512
bf16: True