wisent 0.5.9__tar.gz → 0.5.11__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wisent
-Version: 0.5.9
+Version: 0.5.11
 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.11/wisent/__init__.py

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

{wisent-0.5.9 → wisent-0.5.11}/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
@@ -336,7 +337,7 @@ class WisentModel:
     @torch.inference_mode()
     def generate(
         self,
-        inputs: list[list[ChatMessage]],
+        inputs: list[list[ChatMessage]] | str,
         max_new_tokens: int = 128,
         temperature: float = 0.7,
         top_p: float = 0.95,
@@ -345,6 +346,7 @@ class WisentModel:
         use_steering: bool = False,
         steering_plan: SteeringPlan | None = None,
         enable_thinking: bool = True,
+        prompt_is_formatted: bool = False,
         **gen_kwargs: Any,
     ) -> list[str]:
         """
@@ -352,7 +354,7 @@ class WisentModel:
 
         attributes:
             inputs:
-                list of chat messages (each a list of {'role','content'} dicts).
+                list of chat messages (each a list of {'role','content'} dicts) OR pre-formatted string.
             max_new_tokens:
                 max tokens to generate (beyond the prompt).
             temperature:
@@ -370,6 +372,8 @@ class WisentModel:
                 If None, uses the internal plan.
             enable_thinking:
                 If False, disable thinking/reasoning mode (prevents <think> tags for supported models like Qwen).
+            prompt_is_formatted:
+                If True, inputs is a pre-formatted string with chat template already applied.
             **gen_kwargs:
                 additional kwargs passed to 'model.generate()'.
 
@@ -448,7 +452,22 @@ class WisentModel:
         if use_steering:
             self.apply_steering(steering_plan)
 
-        batch = self._batch_encode(inputs, add_generation_prompt=True, enable_thinking=enable_thinking)
+        if prompt_is_formatted and isinstance(inputs, str):
+            # Direct tokenization of pre-formatted prompt
+            tokenizer_output = self.tokenizer(
+                inputs,
+                return_tensors="pt",
+                padding=False,  # Single prompt, no padding needed
+                truncation=True,  # Avoid errors on long inputs
+                max_length=self.tokenizer.model_max_length  # Use model's actual limit
+            )
+            batch = {
+                "input_ids": tokenizer_output["input_ids"],
+                "attention_mask": tokenizer_output["attention_mask"]
+            }
+        else:
+            # Current behavior: apply chat template
+            batch = self._batch_encode(inputs, add_generation_prompt=True, enable_thinking=enable_thinking)
 
         gen_out = self.hf_model.generate(
             **batch,
@@ -612,7 +631,7 @@ class WisentModel:
     @torch.inference_mode()
     def generate_stream(
         self,
-        inputs: list[list[ChatMessage]],
+        inputs: list[list[ChatMessage]] | str,
         max_new_tokens: int = 128,
         temperature: float = 0.7,
         top_p: float = 0.95,
@@ -622,6 +641,7 @@ class WisentModel:
         skip_prompt: bool = True,
         skip_special_tokens: bool = True,
         enable_thinking: bool = True,
+        prompt_is_formatted: bool = False,
         **gen_kwargs: Any,
     ) -> Iterable[str]:
         """
@@ -630,7 +650,8 @@ class WisentModel:
 
         attributes:
             inputs:
-                list of chat messages (each a list of {'role','content'} dicts). Currently only one conversation is supported.
+                list of chat messages (each a list of {'role','content'} dicts) OR pre-formatted string.
+                Currently only one conversation is supported.
             max_new_tokens:
                 max tokens to generate (beyond the prompt).
             temperature:
@@ -650,6 +671,8 @@ class WisentModel:
                 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).
+            prompt_is_formatted:
+                If True, inputs is a pre-formatted string with chat template already applied.
             **gen_kwargs:
                 additional kwargs passed to 'model.generate()'.
 
@@ -657,14 +680,29 @@ class WisentModel:
             generated text chunks (str), as they become available.
         """
 
-        if len(inputs) != 1:
-            raise ValueError(
-                f"generate_stream currently supports exactly one conversation at a time (got {len(inputs)})."
-            )
         if use_steering:
             self.apply_steering(steering_plan)
 
-        batch = self._batch_encode(inputs, add_generation_prompt=True, enable_thinking=enable_thinking)
+        if prompt_is_formatted and isinstance(inputs, str):
+            # Direct tokenization of pre-formatted prompt
+            tokenizer_output = self.tokenizer(
+                inputs,
+                return_tensors="pt",
+                padding=False,  # Single prompt, no padding needed
+                truncation=True,  # Avoid errors on long inputs
+                max_length=self.tokenizer.model_max_length  # Use model's actual limit
+            )
+            batch = {
+                "input_ids": tokenizer_output["input_ids"],
+                "attention_mask": tokenizer_output["attention_mask"]
+            }
+        else:
+            # Current behavior: apply chat template
+            if not isinstance(inputs, list) or len(inputs) != 1:
+                raise ValueError(
+                    f"generate_stream currently supports exactly one conversation at a time (got {type(inputs)} with {len(inputs) if isinstance(inputs, list) else 'N/A'} items)."
+                )
+            batch = self._batch_encode(inputs, add_generation_prompt=True, enable_thinking=enable_thinking)
 
         streamer = TextIteratorStreamer(
             self.tokenizer,

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

@@ -166,17 +166,19 @@ class MultiSteering:
         max_new_tokens: int = 100,
         temperature: float = 0.7,
         top_p: float = 0.9,
-        enable_thinking: bool = True
+        enable_thinking: bool = True,
+        prompt_is_formatted: bool = False
     ) -> Iterable[str]:
         """Apply the combined steering vector to generate text with streaming.
 
         Args:
             model: WisentModel instance to use for generation
-            prompt: Input prompt
+            prompt: Input prompt (either raw text or pre-formatted with chat template)
             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)
+            prompt_is_formatted: If True, prompt already has chat template applied
 
         Yields:
             Generated text chunks
@@ -191,7 +193,8 @@ class MultiSteering:
             raise MultiSteeringError("No layer information available")
 
         print(f"\n🎯 Applying combined steering vector at layer {self.layer}")
-        print(f"Prompt: {prompt}")
+        print(f"Prompt: {prompt[:100]}..." if len(prompt) > 100 else f"Prompt: {prompt}")
+        print(f"Prompt is formatted: {prompt_is_formatted}")
         print("=" * 50)
 
         # Create SteeringPlan from the combined vector
@@ -202,13 +205,19 @@ class MultiSteering:
             normalize=False  # Already normalized in combine_vectors
         )
 
-        # Format prompt as chat messages
-        messages: list[ChatMessage] = [{"role": "user", "content": prompt}]
+        # Handle prompt formatting
+        if prompt_is_formatted:
+            # Prompt already has chat template applied - pass as string directly
+            inputs = prompt
+        else:
+            # Format prompt as chat messages (current behavior)
+            messages: list[ChatMessage] = [{"role": "user", "content": prompt}]
+            inputs = [messages]
 
         try:
             # Use WisentModel's generate_stream with steering
             yield from model.generate_stream(
-                inputs=[messages],
+                inputs=inputs,
                 max_new_tokens=max_new_tokens,
                 temperature=temperature,
                 top_p=top_p,
@@ -216,7 +225,8 @@ class MultiSteering:
                 steering_plan=steering_plan,
                 skip_prompt=True,
                 skip_special_tokens=True,
-                enable_thinking=enable_thinking
+                enable_thinking=enable_thinking,
+                prompt_is_formatted=prompt_is_formatted
             )
 
         except Exception as e:
@@ -231,17 +241,19 @@ class MultiSteering:
         max_new_tokens: int = 100,
         temperature: float = 0.7,
         top_p: float = 0.9,
-        enable_thinking: bool = True
+        enable_thinking: bool = True,
+        prompt_is_formatted: bool = False
     ) -> str:
         """Apply the combined steering vector to generate text (non-streaming).
 
         Args:
             model: WisentModel instance to use for generation
-            prompt: Input prompt
+            prompt: Input prompt (either raw text or pre-formatted with chat template)
             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)
+            prompt_is_formatted: If True, prompt already has chat template applied
 
         Returns:
             Generated text
@@ -256,7 +268,8 @@ class MultiSteering:
             raise MultiSteeringError("No layer information available")
 
         print(f"\n🎯 Applying combined steering vector at layer {self.layer}")
-        print(f"Prompt: {prompt}")
+        print(f"Prompt: {prompt[:100]}..." if len(prompt) > 100 else f"Prompt: {prompt}")
+        print(f"Prompt is formatted: {prompt_is_formatted}")
         print("=" * 50)
 
         # Create SteeringPlan from the combined vector
@@ -267,19 +280,26 @@ class MultiSteering:
             normalize=False  # Already normalized in combine_vectors
         )
 
-        # Format prompt as chat messages
-        messages: list[ChatMessage] = [{"role": "user", "content": prompt}]
+        # Handle prompt formatting
+        if prompt_is_formatted:
+            # Prompt already has chat template applied - pass as string directly
+            inputs = prompt
+        else:
+            # Format prompt as chat messages (current behavior)
+            messages: list[ChatMessage] = [{"role": "user", "content": prompt}]
+            inputs = [messages]
 
         try:
             # Use WisentModel's generate with steering
             outputs = model.generate(
-                inputs=[messages],
+                inputs=inputs,
                 max_new_tokens=max_new_tokens,
                 temperature=temperature,
                 top_p=top_p,
                 use_steering=True,
                 steering_plan=steering_plan,
-                enable_thinking=enable_thinking
+                enable_thinking=enable_thinking,
+                prompt_is_formatted=prompt_is_formatted
             )
 
             return outputs[0] if outputs else ""

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wisent
-Version: 0.5.9
+Version: 0.5.11
 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

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