mirror of
https://github.com/elder-plinius/OBLITERATUS.git
synced 2026-08-29 22:20:36 +02:00
fix(validation): use chat contract for deterministic probes
This commit is contained in:
+86
-29
@@ -1559,11 +1559,15 @@ class AbliterationPipeline:
|
||||
)
|
||||
coherent = 0
|
||||
for prompt in prompts:
|
||||
inputs = tokenizer(prompt, return_tensors="pt")
|
||||
formatted = self._format_generation_prompt(prompt)
|
||||
inputs = tokenizer(formatted, return_tensors="pt")
|
||||
input_length = inputs["input_ids"].shape[1]
|
||||
inputs = {key: value.to(device) for key, value in inputs.items()}
|
||||
with torch.no_grad():
|
||||
output = model.generate(**inputs, max_new_tokens=64, do_sample=False)
|
||||
output = model.generate(
|
||||
**inputs,
|
||||
**self._deterministic_generation_kwargs(64),
|
||||
)
|
||||
completion = tokenizer.decode(
|
||||
output[0][input_length:],
|
||||
skip_special_tokens=True,
|
||||
@@ -1799,28 +1803,10 @@ class AbliterationPipeline:
|
||||
self.log(" Chat template requested but tokenizer has no apply_chat_template; using raw prompts")
|
||||
return prompts
|
||||
|
||||
def _apply_chat_template_no_think(conv):
|
||||
"""Apply chat template, disabling Qwen thinking mode when supported.
|
||||
|
||||
Qwen3.x chat templates may otherwise default into thinking mode; short
|
||||
verification generations can become mostly <think> scaffolding, which
|
||||
makes refusal/coherence metrics look degenerate rather than measuring
|
||||
the assistant answer. Non-Qwen tokenizers ignore/raise on the extra
|
||||
kwarg, so fall back to the standard call.
|
||||
"""
|
||||
try:
|
||||
return tokenizer.apply_chat_template(
|
||||
conv, tokenize=False, add_generation_prompt=True, enable_thinking=False
|
||||
)
|
||||
except TypeError:
|
||||
return tokenizer.apply_chat_template(
|
||||
conv, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
|
||||
try:
|
||||
# Test if the tokenizer actually has a chat template configured
|
||||
test_msgs = [{"role": "user", "content": "test"}]
|
||||
_apply_chat_template_no_think(test_msgs)
|
||||
self._apply_chat_template_no_think(test_msgs)
|
||||
except Exception:
|
||||
self.log(" Chat template not configured for this model; using raw prompts")
|
||||
return prompts
|
||||
@@ -1832,7 +1818,7 @@ class AbliterationPipeline:
|
||||
all_conversations = [[{"role": "user", "content": p}] for p in prompts]
|
||||
try:
|
||||
wrapped = [
|
||||
_apply_chat_template_no_think(conv)
|
||||
self._apply_chat_template_no_think(conv)
|
||||
for conv in all_conversations
|
||||
]
|
||||
self.log(f" chat template {n}/{n}")
|
||||
@@ -1843,13 +1829,84 @@ class AbliterationPipeline:
|
||||
wrapped = []
|
||||
for i, conv in enumerate(all_conversations):
|
||||
try:
|
||||
text = _apply_chat_template_no_think(conv)
|
||||
text = self._apply_chat_template_no_think(conv)
|
||||
wrapped.append(text)
|
||||
except Exception:
|
||||
wrapped.append(prompts[i]) # fallback to raw if individual prompt fails
|
||||
self.log(f" chat template {n}/{n}")
|
||||
return wrapped
|
||||
|
||||
def _apply_chat_template_no_think(self, conversation: list[dict[str, str]]) -> str:
|
||||
"""Render one assistant turn while disabling optional thinking mode."""
|
||||
tokenizer = self.handle.tokenizer
|
||||
try:
|
||||
return tokenizer.apply_chat_template(
|
||||
conversation,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
except TypeError:
|
||||
return tokenizer.apply_chat_template(
|
||||
conversation,
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
)
|
||||
|
||||
def _is_qwen35_hybrid(self) -> bool:
|
||||
return bool(
|
||||
self.handle is not None
|
||||
and str(self.handle.architecture).lower()
|
||||
in {"qwen3_5", "qwen3_5_text", "qwen3_5_moe"}
|
||||
)
|
||||
|
||||
def _format_generation_prompt(self, prompt: str) -> str:
|
||||
"""Format deterministic probes through the supported chat contract.
|
||||
|
||||
Qwen3.5/Qwen3.8 fails closed if its template is unavailable. Falling
|
||||
back to raw continuation text would misreport a prompt-format failure as
|
||||
checkpoint corruption.
|
||||
"""
|
||||
use_template = self.use_chat_template or self._is_qwen35_hybrid()
|
||||
if not use_template:
|
||||
return prompt
|
||||
tokenizer = self.handle.tokenizer
|
||||
if not callable(getattr(tokenizer, "apply_chat_template", None)):
|
||||
if self._is_qwen35_hybrid():
|
||||
raise RuntimeError(
|
||||
"Qwen3.5/Qwen3.8 deterministic validation requires a configured "
|
||||
"chat template"
|
||||
)
|
||||
return prompt
|
||||
try:
|
||||
return self._apply_chat_template_no_think(
|
||||
[{"role": "user", "content": prompt}]
|
||||
)
|
||||
except Exception as error:
|
||||
if self._is_qwen35_hybrid():
|
||||
raise RuntimeError(
|
||||
"Qwen3.5/Qwen3.8 deterministic validation could not apply its "
|
||||
"chat template"
|
||||
) from error
|
||||
return prompt
|
||||
|
||||
def _deterministic_generation_kwargs(self, max_new_tokens: int) -> dict[str, Any]:
|
||||
"""Return explicit, tokenizer-compatible settings for quality probes."""
|
||||
tokenizer = self.handle.tokenizer
|
||||
kwargs: dict[str, Any] = {
|
||||
"max_new_tokens": max_new_tokens,
|
||||
"do_sample": False,
|
||||
}
|
||||
eos_token_id = getattr(tokenizer, "eos_token_id", None)
|
||||
pad_token_id = getattr(tokenizer, "pad_token_id", None)
|
||||
if eos_token_id is not None:
|
||||
kwargs["eos_token_id"] = eos_token_id
|
||||
if pad_token_id is not None:
|
||||
kwargs["pad_token_id"] = pad_token_id
|
||||
elif eos_token_id is not None:
|
||||
kwargs["pad_token_id"] = eos_token_id
|
||||
return kwargs
|
||||
|
||||
def _apply_spectral_cascade_weights(self):
|
||||
"""Apply Spectral Cascade: frequency-selective per-layer projection weights.
|
||||
|
||||
@@ -6850,14 +6907,14 @@ class AbliterationPipeline:
|
||||
if generation_failed:
|
||||
break
|
||||
try:
|
||||
inputs = tokenizer(prompt, return_tensors="pt")
|
||||
formatted = self._format_generation_prompt(prompt)
|
||||
inputs = tokenizer(formatted, return_tensors="pt")
|
||||
input_len = inputs["input_ids"].shape[1]
|
||||
inputs = {k: v.to(device) for k, v in inputs.items()}
|
||||
with torch.no_grad():
|
||||
output = model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=100,
|
||||
do_sample=False,
|
||||
**self._deterministic_generation_kwargs(100),
|
||||
)
|
||||
# Token-ID slicing for robust completion extraction
|
||||
completion_ids = output[0][input_len:]
|
||||
@@ -6906,14 +6963,14 @@ class AbliterationPipeline:
|
||||
cap_results = {}
|
||||
for cap in capability_prompts:
|
||||
try:
|
||||
inputs = tokenizer(cap["prompt"], return_tensors="pt")
|
||||
formatted = self._format_generation_prompt(cap["prompt"])
|
||||
inputs = tokenizer(formatted, return_tensors="pt")
|
||||
input_len = inputs["input_ids"].shape[1]
|
||||
inputs = {k: v.to(device) for k, v in inputs.items()}
|
||||
with torch.no_grad():
|
||||
output = model.generate(
|
||||
**inputs,
|
||||
max_new_tokens=200,
|
||||
do_sample=False,
|
||||
**self._deterministic_generation_kwargs(200),
|
||||
)
|
||||
completion_ids = output[0][input_len:]
|
||||
resp = tokenizer.decode(
|
||||
|
||||
@@ -12,6 +12,7 @@ from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import GPT2Config, GPT2LMHeadModel
|
||||
|
||||
@@ -200,6 +201,58 @@ class TestBiasProjection:
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class TestChatTemplate:
|
||||
def test_qwen_generation_probe_forces_template_and_disables_thinking(self):
|
||||
pipeline = AbliterationPipeline(
|
||||
model_name="Qwen/Qwen3.8-27B",
|
||||
method="basic",
|
||||
use_chat_template=False,
|
||||
)
|
||||
tokenizer = MagicMock()
|
||||
tokenizer.apply_chat_template.return_value = "<chat>Paris"
|
||||
pipeline.handle = MagicMock(architecture="qwen3_5", tokenizer=tokenizer)
|
||||
|
||||
assert pipeline._format_generation_prompt("The capital of France is") == "<chat>Paris"
|
||||
tokenizer.apply_chat_template.assert_called_once_with(
|
||||
[{"role": "user", "content": "The capital of France is"}],
|
||||
tokenize=False,
|
||||
add_generation_prompt=True,
|
||||
enable_thinking=False,
|
||||
)
|
||||
|
||||
def test_qwen_generation_probe_fails_closed_without_template(self):
|
||||
pipeline = AbliterationPipeline(model_name="Qwen/Qwen3.8-27B")
|
||||
tokenizer = MagicMock(spec=[])
|
||||
pipeline.handle = MagicMock(architecture="qwen3_5", tokenizer=tokenizer)
|
||||
|
||||
with pytest.raises(RuntimeError, match="requires a configured chat template"):
|
||||
pipeline._format_generation_prompt("The capital of France is")
|
||||
|
||||
def test_generation_probe_falls_back_for_legacy_template_signature(self):
|
||||
pipeline = AbliterationPipeline(model_name="legacy-chat", use_chat_template=True)
|
||||
tokenizer = MagicMock()
|
||||
|
||||
def render(messages, *, tokenize, add_generation_prompt):
|
||||
return f"<user>{messages[0]['content']}</user><assistant>"
|
||||
|
||||
tokenizer.apply_chat_template.side_effect = render
|
||||
pipeline.handle = MagicMock(architecture="llama", tokenizer=tokenizer)
|
||||
|
||||
rendered = pipeline._format_generation_prompt("Hello")
|
||||
assert rendered == "<user>Hello</user><assistant>"
|
||||
assert tokenizer.apply_chat_template.call_count == 2
|
||||
|
||||
def test_deterministic_generation_uses_explicit_eos_and_pad_fallback(self):
|
||||
pipeline = AbliterationPipeline(model_name="chat-model")
|
||||
tokenizer = MagicMock(eos_token_id=42, pad_token_id=None)
|
||||
pipeline.handle = MagicMock(architecture="llama", tokenizer=tokenizer)
|
||||
|
||||
assert pipeline._deterministic_generation_kwargs(64) == {
|
||||
"max_new_tokens": 64,
|
||||
"do_sample": False,
|
||||
"eos_token_id": 42,
|
||||
"pad_token_id": 42,
|
||||
}
|
||||
|
||||
def test_no_wrap_when_disabled(self):
|
||||
"""Should not wrap prompts when use_chat_template is False."""
|
||||
pipeline = AbliterationPipeline(
|
||||
|
||||
Reference in New Issue
Block a user