wisent 0.5.8__tar.gz → 0.5.9__tar.gz

{wisent-0.5.8 → wisent-0.5.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wisent
-Version: 0.5.8
+Version: 0.5.9
 Summary: Monitor and guard against harmful content in language models
 Home-page: https://github.com/yourusername/wisent-activation-guardrails
 Author: Wisent Team

wisent-0.5.9/wisent/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.5.9"

{wisent-0.5.8 → wisent-0.5.9}/wisent/core/models/wisent_model.py

@@ -258,6 +258,7 @@ class WisentModel:
         self,
         message: list[ChatMessage],
         add_generation_prompt: bool = True,
+        enable_thinking: bool = True,
     ) -> dict[str, torch.Tensor]:
         """
         Encode a single input in chat format.
@@ -267,6 +268,8 @@ class WisentModel:
                 list of {'role': str, 'content': str} dicts (chat messages).
             add_generation_prompt:
                 If True, append the model's generation prompt at the end.
+            enable_thinking:
+                If False, disable thinking/reasoning mode (prevents <think> tags for supported models like Qwen).
 
         returns:
             dict with 'input_ids' and 'attention_mask' tensors.
@@ -279,10 +282,10 @@ class WisentModel:
             >>> wm._encode_one(msgs, add_generation_prompt=True)
             {"input_ids": tensor([[...]]), "attention_mask": tensor([[...]])}
         """
-      
+
         ids = self.tokenizer.apply_chat_template(
-            message, tokenize=True, add_generation_prompt=add_generation_prompt, return_tensors="pt"
-        )[0] 
+            message, tokenize=True, add_generation_prompt=add_generation_prompt, enable_thinking=enable_thinking, return_tensors="pt"
+        )[0]
         return {
             "input_ids": ids,
             "attention_mask": torch.ones_like(ids),
@@ -292,6 +295,7 @@ class WisentModel:
         self,
         inputs: list[list[ChatMessage]],
         add_generation_prompt: bool = True,
+        enable_thinking: bool = True,
     ) -> dict[str, torch.Tensor]:
         """
         Batch-encode a list of chat messages.
@@ -301,7 +305,9 @@ class WisentModel:
                 list of chat messages (each a list of {'role','content'} dicts).
             add_generation_prompt:
                 If True, append the model's generation prompt at the end of each.
-        
+            enable_thinking:
+                If False, disable thinking/reasoning mode (prevents <think> tags for supported models like Qwen).
+
         returns:
             dict with batched 'input_ids' and 'attention_mask' tensors.
 
@@ -316,10 +322,10 @@ class WisentModel:
             >>> wm._batch_encode([msgs1, msgs2], add_generation_prompt=True)
             {"input_ids": tensor([[...],[...]]), "attention_mask": tensor([[...],[...]])}
         """
-        
+
         singles = []
         for item in inputs:
-            singles.append(self._encode_one(item, add_generation_prompt=add_generation_prompt))
+            singles.append(self._encode_one(item, add_generation_prompt=add_generation_prompt, enable_thinking=enable_thinking))
 
         batch = self.tokenizer.pad(singles, padding=True, return_tensors="pt")
 
@@ -338,6 +344,7 @@ class WisentModel:
         num_return_sequences: int = 1,
         use_steering: bool = False,
         steering_plan: SteeringPlan | None = None,
+        enable_thinking: bool = True,
         **gen_kwargs: Any,
     ) -> list[str]:
         """
@@ -361,6 +368,8 @@ class WisentModel:
             steering_plan:
                 optional SteeringPlan to use for this call only (overrides internal plan).
                 If None, uses the internal plan.
+            enable_thinking:
+                If False, disable thinking/reasoning mode (prevents <think> tags for supported models like Qwen).
             **gen_kwargs:
                 additional kwargs passed to 'model.generate()'.
 
@@ -439,7 +448,7 @@ class WisentModel:
         if use_steering:
             self.apply_steering(steering_plan)
 
-        batch = self._batch_encode(inputs, add_generation_prompt=True)
+        batch = self._batch_encode(inputs, add_generation_prompt=True, enable_thinking=enable_thinking)
 
         gen_out = self.hf_model.generate(
             **batch,
@@ -472,6 +481,7 @@ class WisentModel:
         collect_topk: int = 5,
         use_steering: bool = False,
         steering_plan: SteeringPlan | None = None,
+        enable_thinking: bool = True,
         **gen_kwargs: Any,
     ) -> tuple[list[str], list[GenerationStats]]:
         """
@@ -486,7 +496,7 @@ class WisentModel:
             temperature:
                 sampling temperature (0 = greedy, 1 = default sampling).
             top_p:
-                nucleus sampling probability (0 = no filtering, 1 = full filtering).    
+                nucleus sampling probability (0 = no filtering, 1 = full filtering).
             do_sample:
                 if False, uses greedy decoding (top_k=1).
             num_return_sequences:
@@ -498,6 +508,8 @@ class WisentModel:
             steering_plan:
                 optional SteeringPlan to use for this call only (overrides internal plan).
                 If None, uses the internal plan.
+            enable_thinking:
+                If False, disable thinking/reasoning mode (prevents <think> tags for supported models like Qwen).
             **gen_kwargs:
                 additional kwargs passed to 'model.generate()'.
 
@@ -537,7 +549,7 @@ class WisentModel:
         if use_steering:
             self.apply_steering(steering_plan)
 
-        batch = self._batch_encode(inputs, add_generation_prompt=True)
+        batch = self._batch_encode(inputs, add_generation_prompt=True, enable_thinking=enable_thinking)
 
         out = self.hf_model.generate(
             **batch,
@@ -609,6 +621,7 @@ class WisentModel:
         steering_plan: SteeringPlan | None = None,
         skip_prompt: bool = True,
         skip_special_tokens: bool = True,
+        enable_thinking: bool = True,
         **gen_kwargs: Any,
     ) -> Iterable[str]:
         """
@@ -635,6 +648,8 @@ class WisentModel:
                 if True, the yielded text excludes the input prompt.
             skip_special_tokens:
                 if True, special tokens are removed from the yielded text.
+            enable_thinking:
+                If False, disable thinking/reasoning mode (prevents <think> tags for supported models like Qwen).
             **gen_kwargs:
                 additional kwargs passed to 'model.generate()'.
 
@@ -649,7 +664,7 @@ class WisentModel:
         if use_steering:
             self.apply_steering(steering_plan)
 
-        batch = self._batch_encode(inputs, add_generation_prompt=True)
+        batch = self._batch_encode(inputs, add_generation_prompt=True, enable_thinking=enable_thinking)
 
         streamer = TextIteratorStreamer(
             self.tokenizer,

{wisent-0.5.8 → wisent-0.5.9}/wisent/core/multi_steering.py

@@ -165,7 +165,8 @@ class MultiSteering:
         prompt: str,
         max_new_tokens: int = 100,
         temperature: float = 0.7,
-        top_p: float = 0.9
+        top_p: float = 0.9,
+        enable_thinking: bool = True
     ) -> Iterable[str]:
         """Apply the combined steering vector to generate text with streaming.
 
@@ -175,6 +176,7 @@ class MultiSteering:
             max_new_tokens: Maximum tokens to generate
             temperature: Sampling temperature
             top_p: Top-p sampling parameter
+            enable_thinking: If False, disable thinking/reasoning mode (prevents <think> tags for supported models like Qwen)
 
         Yields:
             Generated text chunks
@@ -213,7 +215,8 @@ class MultiSteering:
                 use_steering=True,
                 steering_plan=steering_plan,
                 skip_prompt=True,
-                skip_special_tokens=True
+                skip_special_tokens=True,
+                enable_thinking=enable_thinking
             )
 
         except Exception as e:
@@ -227,7 +230,8 @@ class MultiSteering:
         prompt: str,
         max_new_tokens: int = 100,
         temperature: float = 0.7,
-        top_p: float = 0.9
+        top_p: float = 0.9,
+        enable_thinking: bool = True
     ) -> str:
         """Apply the combined steering vector to generate text (non-streaming).
 
@@ -237,6 +241,7 @@ class MultiSteering:
             max_new_tokens: Maximum tokens to generate
             temperature: Sampling temperature
             top_p: Top-p sampling parameter
+            enable_thinking: If False, disable thinking/reasoning mode (prevents <think> tags for supported models like Qwen)
 
         Returns:
             Generated text
@@ -273,7 +278,8 @@ class MultiSteering:
                 temperature=temperature,
                 top_p=top_p,
                 use_steering=True,
-                steering_plan=steering_plan
+                steering_plan=steering_plan,
+                enable_thinking=enable_thinking
             )
 
             return outputs[0] if outputs else ""

{wisent-0.5.8 → wisent-0.5.9}/wisent.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wisent
-Version: 0.5.8
+Version: 0.5.9
 Summary: Monitor and guard against harmful content in language models
 Home-page: https://github.com/yourusername/wisent-activation-guardrails
 Author: Wisent Team

wisent-0.5.8/wisent/__init__.py

@@ -1 +0,0 @@
-__version__ = "0.5.8"