Views
No views yet
1class LMWithVectorHead(nn.Module):
2 def __init__(self, model_name, lora_config, output_dim=1):
3 super().__init__()
4 backbone = AutoModel.from_pretrained(model_name, device_map='cpu')
5 # backbone.config.use_cache = False
6 self.peft_model = get_peft_model(backbone, lora_config)
7 self.config = backbone.config
8 hidden_size = backbone.config.hidden_size
9 self.vector_head = nn.Linear(hidden_size, output_dim) # 输出维度为 1
10
11 def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
12 """启用梯度检查点，并处理可能的额外参数"""
13 self.peft_model.enable_input_require_grads()
14 if gradient_checkpointing_kwargs is not None:
15 self.peft_model.gradient_checkpointing_enable(**gradient_checkpointing_kwargs)
16 else:
17 self.peft_model.gradient_checkpointing_enable()
18
19 def forward(self, input_ids, attention_mask=None, labels=None):
20 # if hasattr(self.peft_model, "gradient_checkpointing"):
21 # print(f"✅ 梯度检查点已启用 - 当前模式: {self.peft_model.is_gradient_checkpointing}")
22 # else:
23 # print("❌ 梯度检查点未正确初始化")
24 outputs = self.peft_model(
25 input_ids=input_ids,
26 attention_mask=attention_mask,
27 return_dict=True
28 )
29 # 获取最后一个 token 的隐藏状态
30 last_hidden = outputs.last_hidden_state # [B, T, H]
31 cls_hidden = last_hidden[:, -1, :] # [B, H]
32 logits = self.vector_head(cls_hidden) # [B, 1]
33 logits = torch.sigmoid(logits).squeeze(-1) # 添加 sigmoid 并压缩至 [B]
34
35 loss = None
36 if labels is not None:
37 loss_fct = nn.MSELoss() # 使用 MSE 损失
38 loss = loss_fct(logits, labels) # 计算 logits 和 labels 的 MSE
39
40 return CausalLMOutput(
41 loss=loss,
42 logits=logits
43 )1base_model = AutoModel.from_pretrained(args.model_path)
2
3 # 2. 加载训练好的LoRA适配器到基础模型上
4 peft_model = PeftModel.from_pretrained(
5 base_model, # 使用基础模型，而不是model.peft_model
6 args.lora_path,
7 adapter_name="default"
8 )
9
10 # 3. 创建完整模型结构
11 lora_config = LoraConfig(
12 r=args.r,
13 lora_alpha=args.alpha,
14 target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
15 "gate_proj", "up_proj", "down_proj"],
16 lora_dropout=args.lora_dropout,
17 bias="none",
18 )
19 model = LMWithVectorHead(args.model_path, lora_config)
20
21 # 4. 替换为已加载LoRA的模型
22 model.peft_model = peft_model
23
24 # 5. 加载分类头权重
25 state_dict = torch.load(args.vector_head_path, map_location=device)
26 model.vector_head.load_state_dict(state_dict)
27
28 # 6. 激活适配器并移动到设备
29 model.peft_model.set_adapter("default")
30 model = model.to(device)
31
32 # 评估模式
33 model.eval()1/lora
2├── greedy_answer_conf
3│ └── long_qa
4│ └── batchsize16_accumulation8_epochs10_weightdecay0.1_r8_alpha16_loradropout0.0 (training configuration)
5│ ├── best_checkpoints
6│ │ ├── lora_epoch_best/ # Path to LoRA module
7│ │ └── vector_head_epoch_best.pt # Path to Linear Head weights
8│ └── test_losses.json # Test loss for each epoch
9│
10├── hybrid_answer_conf
11│ └── long_qa
12│ ├── batchsize16_accumulation8_epochs10_weightdecay0.1_r8_alpha16_loradropout0.0 (560k samples)
13│ ├── batchsize16_accumulation8_epochs50_weightdecay0.1_r8_alpha16_loradropout0.0_1k_training_samples (1k samples)
14│ └── batchsize16_accumulation8_epochs50_weightdecay0.1_r8_alpha16_loradropout0.0_2k_training_samples (2k samples)
15│
16└── right_answer_conf
17 └── long_qa
18 └── ... # Same format as above
19
20/mlp
21...bash run_eval.sh| Column | Meaning |
|---|---|
| N-Prob | Normalized probability baseline |
| Cons-Sem | Consistency with semantic agreement |
| Eli-Only | Elicitation-only model |
| Cal-Only | Calibration-only model |
| EliCal | Full EliCal method |
_mlp at the end of a script name, it means that only a classification head is added to the model without using LoRA. This is not the main focus of the paper, but simply an ablation study.We provide all the plotting functions used in the paper inhonesty_alignment/draw.py.For more details, please refer to our paper.