Merge pull request #169 from jmagly/fix/168-qwen38-chat-baseline

fix(validation): honor Qwen chat contract in quality probes

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 <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(

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 = "<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(