wisent 0.5.8__py3-none-any.whl → 0.5.10__py3-none-any.whl

wisent/__init__.py

@@ -1 +1 @@
-__version__ = "0.5.8"
+__version__ = "0.5.10"

wisent/core/models/wisent_model.py

@@ -95,6 +95,7 @@ class WisentModel:
         elif self.device == "cuda":
             load_kwargs["dtype"] = torch.float16
             load_kwargs["device_map"] = "auto"
+            load_kwargs["attn_implementation"] = "flash_attention_2"  # Use flash attention for CUDA
         else:
             load_kwargs["dtype"] = torch.float32
             load_kwargs["device_map"] = None
@@ -258,6 +259,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 +269,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 +283,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 +296,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 +306,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 +323,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 +345,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 +369,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 +449,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 +482,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 +497,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 +509,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 +550,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 +622,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 +649,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 +665,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/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.dist-info → wisent-0.5.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wisent
-Version: 0.5.8
+Version: 0.5.10
 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.dist-info → wisent-0.5.10.dist-info}/RECORD

@@ -1,4 +1,4 @@
-wisent/__init__.py,sha256=bDuZ37zImJZsQ3a4pW87q4kg-zsIBrUFAv1aumIf_7k,22
+wisent/__init__.py,sha256=1nlPInsRzDbcDPveZ3ghSJ6v6KveN9n6gnj-twW4DkI,23
 wisent/benchmarks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 wisent/benchmarks/coding/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 wisent/benchmarks/coding/metrics/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -64,7 +64,7 @@ wisent/core/managed_cached_benchmarks.py,sha256=JbvpZ1fgSuQQhyQVKEvqrQZRHGqfnjo9
 wisent/core/mixed_benchmark_sampler.py,sha256=tKQCHUXVuYeCyx4VZt8O1hGyB-TOY_SQ_SYi8cyApII,13585
 wisent/core/model_config_manager.py,sha256=rQAdSmk3GFlZXyHp3fSV1bORxiZWhmzIz1uo3H4JtkA,12009
 wisent/core/model_persistence.py,sha256=6_vc1Ndujd4v0O68giINSTvYhmb7-AiacWwAbqLOrls,10636
-wisent/core/multi_steering.py,sha256=o4YzkEMeOk8rWXC4JPa3gwVZi4_CjjwyLGXlkWxpPOw,11869
+wisent/core/multi_steering.py,sha256=EMaKn4dZPlAsFupEUQZlxTZGJ0-ofpLcTCKQk8HaZL8,12295
 wisent/core/parser.py,sha256=_YDeSuQMx0zNknz9rX3Ls1YPT1x5eohoY8rfjeoqxV8,69091
 wisent/core/representation.py,sha256=hBl_N9qbr5Gsa7GCQ0nMWRm82RqYEfhd9cyf0PPH5LY,195
 wisent/core/sample_size_optimizer.py,sha256=6wegGXZpdGpiR4R0YJ1D2JqLr6yinMndEx2gB5FL80s,23666
@@ -134,7 +134,7 @@ wisent/core/evaluators/oracles/interactive.py,sha256=f3v2_N17fKzGyeOxONRJbrbn8i5
 wisent/core/evaluators/oracles/nlp_evaluator.py,sha256=KxbnF-I2IFbBQpoYyjQKGbYh4NErsEuhTCRYX_Tob8o,18220
 wisent/core/evaluators/oracles/user_specified.py,sha256=V1dKrNj3Oq7UC_I7DT0WGnktP7R_DSW6UAwDdrA8SnE,2360
 wisent/core/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-wisent/core/models/wisent_model.py,sha256=yJBcz3GjR7O-ySTV2vvOsOrL9xDvXsG0W9Gr0HR_0sc,28729
+wisent/core/models/wisent_model.py,sha256=-QJRrPxQPduDyjH0l9PDZC0cdoBzyrQQ_bgeImfGwfI,29873
 wisent/core/models/core/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 wisent/core/models/core/atoms.py,sha256=_Bpz0Sfiq6_VswThIltUwNGj_ukl5MhAg8RrgMKwEBM,15756
 wisent/core/optuna/__init__.py,sha256=sTfwRnrRyKrCNVsF_qCjBDFEZC0ZmUZ7m6IE0iHfTVs,1914
@@ -213,8 +213,8 @@ wisent/synthetic/generators/diversities/core/__init__.py,sha256=47DEQpj8HBSa-_TI
 wisent/synthetic/generators/diversities/core/core.py,sha256=TjSj5T7NE5kRH-ABcFqb1Hz_j3Z6F_TcV-95uHD5Xw8,2201
 wisent/synthetic/generators/diversities/methods/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 wisent/synthetic/generators/diversities/methods/fast_diversity.py,sha256=Z2UzTbzyJFM_ToxCoXM_LQQQ1Jc6BZknrbpikTG1MRw,8522
-wisent-0.5.8.dist-info/licenses/LICENSE,sha256=wy0iaw8b2tyqZAfKHib3lP3PJ9o88FDCg92oUHh3sDQ,1073
-wisent-0.5.8.dist-info/METADATA,sha256=u406l73QL6jE0Jw5GZW3O4RKZkgy-NTs_DOgLIAoGDM,2424
-wisent-0.5.8.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-wisent-0.5.8.dist-info/top_level.txt,sha256=2Ts9Iyldnb3auIN2HBBaHPknRy7nSRDm2f6RGzYgr8A,7
-wisent-0.5.8.dist-info/RECORD,,
+wisent-0.5.10.dist-info/licenses/LICENSE,sha256=wy0iaw8b2tyqZAfKHib3lP3PJ9o88FDCg92oUHh3sDQ,1073
+wisent-0.5.10.dist-info/METADATA,sha256=-CwqNc9Sz5brGHuNi1lMz34YNVqKNHRft4QVlOAjZvM,2425
+wisent-0.5.10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+wisent-0.5.10.dist-info/top_level.txt,sha256=2Ts9Iyldnb3auIN2HBBaHPknRy7nSRDm2f6RGzYgr8A,7
+wisent-0.5.10.dist-info/RECORD,,