diff --git a/obliteratus/abliterate.py b/obliteratus/abliterate.py index 41ace51..365047b 100644 --- a/obliteratus/abliterate.py +++ b/obliteratus/abliterate.py @@ -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 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( diff --git a/tests/test_abliterate_extended.py b/tests/test_abliterate_extended.py index c80162c..33c7a5c 100644 --- a/tests/test_abliterate_extended.py +++ b/tests/test_abliterate_extended.py @@ -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 = "Paris" + pipeline.handle = MagicMock(architecture="qwen3_5", tokenizer=tokenizer) + + assert pipeline._format_generation_prompt("The capital of France is") == "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"{messages[0]['content']}" + + tokenizer.apply_chat_template.side_effect = render + pipeline.handle = MagicMock(architecture="llama", tokenizer=tokenizer) + + rendered = pipeline._format_generation_prompt("Hello") + assert rendered == "Hello" + 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(