jfpio commited on
Commit
95f1c8b
·
verified ·
1 Parent(s): dfc76e6

Upload folder using huggingface_hub

Browse files
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ tokenizer.json filter=lfs diff=lfs merge=lfs -text
README.md ADDED
@@ -0,0 +1,183 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Wisent-Qwen2.5-Coder-7B-Instruct with CAA Steering
2
+
3
+ ## Model Description
4
+
5
+ This is an enhanced version of Qwen2.5-Coder-7B-Instruct that integrates **Contrastive Activation Addition (CAA)** steering directly into the model architecture. The steering parameters have been optimized using Optuna to improve code generation quality on the MBPP Plus benchmark.
6
+
7
+ ### Key Features
8
+
9
+ - 🚀 **Automatic CAA Steering**: No manual hook management required
10
+ - 🎯 **Optimized Parameters**: Layer 24, α=0.9
11
+ - 🗂️ **Trait-Based Organization**: Steering vectors organized by traits
12
+ - 🔧 **Runtime Configurable**: Adjust or disable steering on the fly
13
+ - 🤗 **HuggingFace Compatible**: Works with standard transformers API
14
+
15
+ ## Installation
16
+
17
+ ```bash
18
+ pip install transformers torch
19
+ ```
20
+
21
+ ## Quick Start
22
+
23
+ ```python
24
+ from transformers import AutoModelForCausalLM, AutoTokenizer
25
+
26
+ # Load model - CAA steering is automatically applied!
27
+ model = AutoModelForCausalLM.from_pretrained("./huggingface_qwen_generated", trust_remote_code=True)
28
+ tokenizer = AutoTokenizer.from_pretrained("./huggingface_qwen_generated")
29
+
30
+ # Generate code
31
+ prompt = "Write a Python function to calculate the factorial of a number"
32
+ inputs = tokenizer(prompt, return_tensors="pt")
33
+ outputs = model.generate(**inputs, max_new_tokens=256, temperature=0.2)
34
+ response = tokenizer.decode(outputs[0], skip_special_tokens=True)
35
+ print(response)
36
+ ```
37
+
38
+ ## Advanced Usage
39
+
40
+ ### Adjusting Steering Strength
41
+
42
+ ```python
43
+ # Increase steering strength for stronger safety alignment
44
+ model.set_caa_alpha(1.2)
45
+
46
+ # Decrease for more creative outputs
47
+ model.set_caa_alpha(0.5)
48
+ ```
49
+
50
+ ### Disabling CAA Steering
51
+
52
+ ```python
53
+ # Disable CAA to get baseline model behavior
54
+ model.set_caa_enabled(False)
55
+
56
+ # Re-enable CAA
57
+ model.set_caa_enabled(True)
58
+ ```
59
+
60
+ ### Accessing Steering Configuration
61
+
62
+ ```python
63
+ print(f"CAA Layer: {model.caa_layer_id}")
64
+ print(f"CAA Alpha: {model.caa_alpha}")
65
+ print(f"Steering Method: {model.steering_method}")
66
+ ```
67
+
68
+ ### Trait-Based Vector Organization
69
+
70
+ The model uses a trait-based organization for steering vectors:
71
+
72
+ ```
73
+ vectors/
74
+ ├── coding/ # Current: Optimized for code generation
75
+ ├── safety/ # Future: Safety-aligned behavior
76
+ ├── creativity/ # Future: Enhanced creative outputs
77
+ ├── helpfulness/ # Future: Improved helpfulness
78
+ └── reasoning/ # Future: Enhanced logical reasoning
79
+ ```
80
+
81
+ To switch traits, simply update the configuration:
82
+
83
+ ```json
84
+ {
85
+ "steering_vector_path": "./vectors/safety/steering_vector.safetensors"
86
+ }
87
+ ```
88
+
89
+ ## Technical Details
90
+
91
+ ### CAA Steering Parameters
92
+
93
+ - **Steering Method**: Contrastive Activation Addition (CAA)
94
+ - **Optimal Layer**: 24 (out of 28 transformer layers)
95
+ - **Steering Strength (α)**: 0.9
96
+ - **Vector Format**: Safetensors format for efficient loading and HuggingFace compatibility
97
+ - **Vector Dimension**: 3584 (pre-normalized during training)
98
+ - **Storage Path**: `./vectors/coding/steering_vector.safetensors`
99
+
100
+ ### How It Works
101
+
102
+ 1. **Trait-based Organization**: Steering vectors are organized by behavioral traits (`vectors/{trait}/`)
103
+ 2. **Dynamic Loading**: The model loads the specified steering vector from the configured path
104
+ 3. **Layer Application**: Steering is applied to hidden states at layer 24 during forward pass
105
+ 4. **Generation Integration**: Steering affects the last token position during generation
106
+ 5. **Configurable Strength**: The α parameter (default: 0.9) controls steering intensity
107
+ 6. **Pre-optimized Vectors**: Steering vectors are pre-normalized and ready for immediate use
108
+
109
+ ### Optimization Process
110
+
111
+ The CAA parameters were optimized using:
112
+ - **Framework**: Optuna with TPE sampler
113
+ - **Search Space**: Layers 15-28, α ∈ [0.1, 5.0]
114
+ - **Objective**: Maximize accuracy on MBPP Plus validation set
115
+ - **Best Validation Score**: 64% accuracy
116
+
117
+ ## Model Architecture
118
+
119
+ ```
120
+ WisentQwen2ForCausalLM
121
+ ├── Base: Qwen2.5-Coder-7B-Instruct
122
+ ├── CAA Integration: Layer 24
123
+ ├── Steering Vector: ./vectors/coding/steering_vector.safetensors
124
+ └── Auto-applied during generation
125
+ ```
126
+
127
+ ## File Structure
128
+
129
+ ```
130
+ huggingface_qwen_generated/
131
+ ├── config.json # Model configuration with CAA params
132
+ ├── modeling_wisent_qwen.py # Custom model class
133
+ ├── tokenizer files # Standard Qwen tokenizer
134
+ ├── wisent_config.json # Optimization results
135
+ └── vectors/ # Trait-based steering vectors
136
+ └── coding/
137
+ └── steering_vector.safetensors # Optimized coding steering vector
138
+ ```
139
+
140
+ ## Evaluation
141
+
142
+ ### MBPP Plus Benchmark
143
+
144
+ The model should be evaluated on the complete MBPP Plus dataset (378 problems) to measure improvement over the baseline. Expected improvements based on validation results.
145
+
146
+ ### Running Evaluation
147
+
148
+ ```python
149
+ # Use with bigcode-evaluation-harness
150
+ from transformers import AutoModelForCausalLM
151
+
152
+ model = AutoModelForCausalLM.from_pretrained(
153
+ "./huggingface_qwen_generated",
154
+ trust_remote_code=True
155
+ )
156
+
157
+ # CAA steering is automatically applied during evaluation!
158
+ # No manual hooks or modifications needed
159
+ ```
160
+
161
+ ## Citation
162
+
163
+ If you use this model, please cite:
164
+
165
+ ```bibtex
166
+ @software{wisent_qwen_caa_2025,
167
+ title={Wisent-Qwen2.5-Coder with CAA Steering},
168
+ author={Wisent AI},
169
+ year={2025},
170
+ url={https://github.com/wisent-ai/wisent-guard}
171
+ }
172
+ ```
173
+
174
+ ## License
175
+
176
+ This model inherits the license from the base Qwen2.5-Coder-7B-Instruct model. Please refer to the original model's license for usage terms.
177
+
178
+ ## Acknowledgments
179
+
180
+ - Base model: Qwen2.5-Coder-7B-Instruct by Alibaba
181
+ - CAA method: Contrastive Activation Addition
182
+ - Optimization: Optuna framework
183
+ - Implementation: Wisent Guard framework
added_tokens.json ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "</tool_call>": 151658,
3
+ "<tool_call>": 151657,
4
+ "<|box_end|>": 151649,
5
+ "<|box_start|>": 151648,
6
+ "<|endoftext|>": 151643,
7
+ "<|file_sep|>": 151664,
8
+ "<|fim_middle|>": 151660,
9
+ "<|fim_pad|>": 151662,
10
+ "<|fim_prefix|>": 151659,
11
+ "<|fim_suffix|>": 151661,
12
+ "<|im_end|>": 151645,
13
+ "<|im_start|>": 151644,
14
+ "<|image_pad|>": 151655,
15
+ "<|object_ref_end|>": 151647,
16
+ "<|object_ref_start|>": 151646,
17
+ "<|quad_end|>": 151651,
18
+ "<|quad_start|>": 151650,
19
+ "<|repo_name|>": 151663,
20
+ "<|video_pad|>": 151656,
21
+ "<|vision_end|>": 151653,
22
+ "<|vision_pad|>": 151654,
23
+ "<|vision_start|>": 151652
24
+ }
chat_template.jinja ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {%- if tools %}
2
+ {{- '<|im_start|>system\n' }}
3
+ {%- if messages[0]['role'] == 'system' %}
4
+ {{- messages[0]['content'] }}
5
+ {%- else %}
6
+ {{- 'You are Qwen, created by Alibaba Cloud. You are a helpful assistant.' }}
7
+ {%- endif %}
8
+ {{- "\n\n# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
9
+ {%- for tool in tools %}
10
+ {{- "\n" }}
11
+ {{- tool | tojson }}
12
+ {%- endfor %}
13
+ {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
14
+ {%- else %}
15
+ {%- if messages[0]['role'] == 'system' %}
16
+ {{- '<|im_start|>system\n' + messages[0]['content'] + '<|im_end|>\n' }}
17
+ {%- else %}
18
+ {{- '<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n' }}
19
+ {%- endif %}
20
+ {%- endif %}
21
+ {%- for message in messages %}
22
+ {%- if (message.role == "user") or (message.role == "system" and not loop.first) or (message.role == "assistant" and not message.tool_calls) %}
23
+ {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
24
+ {%- elif message.role == "assistant" %}
25
+ {{- '<|im_start|>' + message.role }}
26
+ {%- if message.content %}
27
+ {{- '\n' + message.content }}
28
+ {%- endif %}
29
+ {%- for tool_call in message.tool_calls %}
30
+ {%- if tool_call.function is defined %}
31
+ {%- set tool_call = tool_call.function %}
32
+ {%- endif %}
33
+ {{- '\n<tool_call>\n{"name": "' }}
34
+ {{- tool_call.name }}
35
+ {{- '", "arguments": ' }}
36
+ {{- tool_call.arguments | tojson }}
37
+ {{- '}\n</tool_call>' }}
38
+ {%- endfor %}
39
+ {{- '<|im_end|>\n' }}
40
+ {%- elif message.role == "tool" %}
41
+ {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != "tool") %}
42
+ {{- '<|im_start|>user' }}
43
+ {%- endif %}
44
+ {{- '\n<tool_response>\n' }}
45
+ {{- message.content }}
46
+ {{- '\n</tool_response>' }}
47
+ {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
48
+ {{- '<|im_end|>\n' }}
49
+ {%- endif %}
50
+ {%- endif %}
51
+ {%- endfor %}
52
+ {%- if add_generation_prompt %}
53
+ {{- '<|im_start|>assistant\n' }}
54
+ {%- endif %}
config.json ADDED
@@ -0,0 +1,127 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "vocab_size": 152064,
3
+ "max_position_embeddings": 32768,
4
+ "hidden_size": 3584,
5
+ "intermediate_size": 18944,
6
+ "num_hidden_layers": 28,
7
+ "num_attention_heads": 28,
8
+ "use_sliding_window": false,
9
+ "sliding_window": null,
10
+ "max_window_layers": 28,
11
+ "num_key_value_heads": 4,
12
+ "hidden_act": "silu",
13
+ "initializer_range": 0.02,
14
+ "rms_norm_eps": 1e-06,
15
+ "use_cache": true,
16
+ "rope_theta": 1000000.0,
17
+ "rope_scaling": null,
18
+ "attention_dropout": 0.0,
19
+ "layer_types": [
20
+ "full_attention",
21
+ "full_attention",
22
+ "full_attention",
23
+ "full_attention",
24
+ "full_attention",
25
+ "full_attention",
26
+ "full_attention",
27
+ "full_attention",
28
+ "full_attention",
29
+ "full_attention",
30
+ "full_attention",
31
+ "full_attention",
32
+ "full_attention",
33
+ "full_attention",
34
+ "full_attention",
35
+ "full_attention",
36
+ "full_attention",
37
+ "full_attention",
38
+ "full_attention",
39
+ "full_attention",
40
+ "full_attention",
41
+ "full_attention",
42
+ "full_attention",
43
+ "full_attention",
44
+ "full_attention",
45
+ "full_attention",
46
+ "full_attention",
47
+ "full_attention"
48
+ ],
49
+ "return_dict": true,
50
+ "output_hidden_states": false,
51
+ "torchscript": false,
52
+ "torch_dtype": "bfloat16",
53
+ "use_bfloat16": false,
54
+ "tf_legacy_loss": false,
55
+ "pruned_heads": {},
56
+ "tie_word_embeddings": false,
57
+ "chunk_size_feed_forward": 0,
58
+ "is_encoder_decoder": false,
59
+ "is_decoder": false,
60
+ "cross_attention_hidden_size": null,
61
+ "add_cross_attention": false,
62
+ "tie_encoder_decoder": false,
63
+ "max_length": 20,
64
+ "min_length": 0,
65
+ "do_sample": false,
66
+ "early_stopping": false,
67
+ "num_beams": 1,
68
+ "num_beam_groups": 1,
69
+ "diversity_penalty": 0.0,
70
+ "temperature": 1.0,
71
+ "top_k": 50,
72
+ "top_p": 1.0,
73
+ "typical_p": 1.0,
74
+ "repetition_penalty": 1.0,
75
+ "length_penalty": 1.0,
76
+ "no_repeat_ngram_size": 0,
77
+ "encoder_no_repeat_ngram_size": 0,
78
+ "bad_words_ids": null,
79
+ "num_return_sequences": 1,
80
+ "output_scores": false,
81
+ "return_dict_in_generate": false,
82
+ "forced_bos_token_id": null,
83
+ "forced_eos_token_id": null,
84
+ "remove_invalid_values": false,
85
+ "exponential_decay_length_penalty": null,
86
+ "suppress_tokens": null,
87
+ "begin_suppress_tokens": null,
88
+ "architectures": [
89
+ "WisentQwen2ForCausalLM"
90
+ ],
91
+ "finetuning_task": null,
92
+ "id2label": {
93
+ "0": "LABEL_0",
94
+ "1": "LABEL_1"
95
+ },
96
+ "label2id": {
97
+ "LABEL_0": 0,
98
+ "LABEL_1": 1
99
+ },
100
+ "tokenizer_class": null,
101
+ "prefix": null,
102
+ "bos_token_id": 151643,
103
+ "pad_token_id": null,
104
+ "eos_token_id": 151645,
105
+ "sep_token_id": null,
106
+ "decoder_start_token_id": null,
107
+ "task_specific_params": null,
108
+ "problem_type": null,
109
+ "_name_or_path": "Qwen/Qwen2.5-Coder-7B-Instruct",
110
+ "transformers_version": "4.53.3",
111
+ "model_type": "wisent_qwen2",
112
+ "output_attentions": false,
113
+ "auto_map": {
114
+ "AutoConfig": "modeling_wisent_qwen.WisentQwen2Config",
115
+ "AutoModelForCausalLM": "modeling_wisent_qwen.WisentQwen2ForCausalLM"
116
+ },
117
+ "caa_enabled": true,
118
+ "caa_layer_id": 24,
119
+ "caa_alpha": 0.9,
120
+ "steering_method": "caa",
121
+ "wisent_optimization": {
122
+ "best_value": 0.64,
123
+ "timestamp": "20250818_221712",
124
+ "commit_hash": "a2181df6155f0d5d20170f307b61d10e74d31889"
125
+ },
126
+ "steering_vector_path": "./vectors/coding/steering_vector.safetensors"
127
+ }
merges.txt ADDED
The diff for this file is too large to render. See raw diff
 
modeling_wisent_qwen.py ADDED
@@ -0,0 +1,295 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Wisent-enhanced Qwen2 model with integrated CAA (Contrastive Activation Addition) steering.
3
+
4
+ This model automatically applies CAA steering during generation without requiring manual hooks.
5
+ The steering parameters are optimized using Optuna and stored in the model configuration.
6
+ """
7
+
8
+ from typing import Optional, Tuple, Union, List
9
+ import torch
10
+ import torch.nn as nn
11
+ from transformers import Qwen2ForCausalLM, Qwen2Config
12
+ from transformers.modeling_outputs import CausalLMOutputWithPast
13
+ from transformers.cache_utils import Cache
14
+
15
+
16
+ class WisentQwen2Config(Qwen2Config):
17
+ """Extended Qwen2 configuration with CAA steering parameters."""
18
+
19
+ model_type = "wisent_qwen2"
20
+
21
+ def __init__(
22
+ self,
23
+ caa_enabled: bool = True,
24
+ caa_layer_id: int = 24,
25
+ caa_alpha: float = 0.9,
26
+ steering_vector_path: str = "./vectors/coding/steering_vector.safetensors",
27
+ steering_method: str = "caa",
28
+ **kwargs
29
+ ):
30
+ super().__init__(**kwargs)
31
+ self.caa_enabled = caa_enabled
32
+ self.caa_layer_id = caa_layer_id
33
+ self.caa_alpha = caa_alpha
34
+ self.steering_vector_path = steering_vector_path
35
+ self.steering_method = steering_method
36
+
37
+
38
+ class WisentQwen2ForCausalLM(Qwen2ForCausalLM):
39
+ """
40
+ Qwen2 model with integrated CAA steering for improved code generation.
41
+
42
+ This model automatically applies Contrastive Activation Addition (CAA) steering
43
+ during the forward pass, eliminating the need for manual hook management.
44
+ """
45
+
46
+ config_class = WisentQwen2Config
47
+
48
+ def __init__(self, config: WisentQwen2Config):
49
+ super().__init__(config)
50
+
51
+ # CAA steering parameters
52
+ self.caa_enabled = config.caa_enabled
53
+ self.caa_layer_id = config.caa_layer_id
54
+ self.caa_alpha = config.caa_alpha
55
+ self.steering_method = config.steering_method
56
+
57
+ # Load steering vector from file
58
+ self.steering_vector = None
59
+ if self.caa_enabled:
60
+ self._load_steering_vector_from_file(config.steering_vector_path)
61
+
62
+ # Hook handle for cleanup
63
+ self._steering_hook_handle = None
64
+
65
+ def _load_steering_vector_from_file(self, path: str):
66
+ """Load the CAA steering vector from safetensors or pytorch file."""
67
+ import os
68
+ try:
69
+ # Try relative path first
70
+ if os.path.exists(path):
71
+ vector_path = path
72
+ # Try path relative to model directory
73
+ elif os.path.exists(os.path.join(os.path.dirname(__file__), path)):
74
+ vector_path = os.path.join(os.path.dirname(__file__), path)
75
+ else:
76
+ print(f"Warning: Steering vector not found at {path}, CAA disabled")
77
+ self.caa_enabled = False
78
+ return
79
+
80
+ # Load based on file extension
81
+ if vector_path.endswith('.safetensors'):
82
+ # Load from safetensors format (preferred)
83
+ try:
84
+ from safetensors.torch import load_file
85
+ steering_data = load_file(vector_path)
86
+ self.steering_vector = steering_data['steering_vector']
87
+ except ImportError:
88
+ print("Warning: safetensors not installed, install with: pip install safetensors")
89
+ self.caa_enabled = False
90
+ return
91
+ else:
92
+ # Load from pytorch format (fallback)
93
+ steering_data = torch.load(vector_path, map_location='cpu')
94
+
95
+ # Handle different storage formats
96
+ if isinstance(steering_data, dict):
97
+ if 'vector' in steering_data:
98
+ self.steering_vector = steering_data['vector']
99
+ elif 'steering_vector' in steering_data:
100
+ self.steering_vector = steering_data['steering_vector']
101
+ else:
102
+ # Assume the dict values are the vectors
103
+ self.steering_vector = next(iter(steering_data.values()))
104
+ else:
105
+ self.steering_vector = steering_data
106
+
107
+ # Ensure it's a tensor
108
+ if not isinstance(self.steering_vector, torch.Tensor):
109
+ self.steering_vector = torch.tensor(self.steering_vector)
110
+
111
+ print(f"✅ Loaded CAA steering vector from {vector_path}: shape {self.steering_vector.shape}, norm {torch.norm(self.steering_vector).item():.4f}")
112
+
113
+ except Exception as e:
114
+ print(f"Warning: Failed to load steering vector: {e}, CAA disabled")
115
+ self.caa_enabled = False
116
+ self.steering_vector = None
117
+
118
+ def _apply_caa_steering(self, module, input, output):
119
+ """
120
+ Hook function that applies CAA steering to the specified layer.
121
+
122
+ This follows the implementation from wisent_guard/core/steering_methods/caa.py
123
+ and the patterns from wisent_guard/core/optuna/optuna_pipeline.py
124
+ """
125
+ if not self.caa_enabled or self.steering_vector is None:
126
+ return output
127
+
128
+ # Extract hidden states from output
129
+ if isinstance(output, tuple):
130
+ hidden_states = output[0]
131
+ else:
132
+ hidden_states = output
133
+
134
+ # Apply steering to the last token position (standard CAA behavior)
135
+ # This matches the implementation in optuna_pipeline.py lines 744-746
136
+ if hidden_states.dim() == 3: # [batch, seq, hidden]
137
+ # Move steering vector to the same device and dtype
138
+ steering_vector = self.steering_vector.to(hidden_states.device, hidden_states.dtype)
139
+
140
+ # Apply steering with configured alpha (strength)
141
+ # Steering is applied to the last token position
142
+ hidden_states[:, -1:, :] = hidden_states[:, -1:, :] + self.caa_alpha * steering_vector.unsqueeze(0).unsqueeze(0)
143
+
144
+ # Return modified output
145
+ if isinstance(output, tuple):
146
+ return (hidden_states,) + output[1:]
147
+ else:
148
+ return hidden_states
149
+
150
+ def forward(
151
+ self,
152
+ input_ids: torch.LongTensor = None,
153
+ attention_mask: Optional[torch.Tensor] = None,
154
+ position_ids: Optional[torch.LongTensor] = None,
155
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
156
+ inputs_embeds: Optional[torch.FloatTensor] = None,
157
+ labels: Optional[torch.LongTensor] = None,
158
+ use_cache: Optional[bool] = None,
159
+ output_attentions: Optional[bool] = None,
160
+ output_hidden_states: Optional[bool] = None,
161
+ return_dict: Optional[bool] = None,
162
+ cache_position: Optional[torch.LongTensor] = None,
163
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
164
+ """
165
+ Forward pass with automatic CAA steering application.
166
+
167
+ The steering is applied via a forward hook on the specified layer,
168
+ following the pattern from optuna_pipeline.py.
169
+ """
170
+
171
+ # Register CAA steering hook if enabled and not already registered
172
+ if self.caa_enabled and self.steering_vector is not None and self._steering_hook_handle is None:
173
+ target_layer = self.model.layers[self.caa_layer_id]
174
+ self._steering_hook_handle = target_layer.register_forward_hook(self._apply_caa_steering)
175
+
176
+ # Call parent forward method
177
+ outputs = super().forward(
178
+ input_ids=input_ids,
179
+ attention_mask=attention_mask,
180
+ position_ids=position_ids,
181
+ past_key_values=past_key_values,
182
+ inputs_embeds=inputs_embeds,
183
+ labels=labels,
184
+ use_cache=use_cache,
185
+ output_attentions=output_attentions,
186
+ output_hidden_states=output_hidden_states,
187
+ return_dict=return_dict,
188
+ cache_position=cache_position if hasattr(self, 'cache_position') else None,
189
+ )
190
+
191
+ return outputs
192
+
193
+ def generate(self, *args, **kwargs):
194
+ """
195
+ Generate method with automatic CAA steering.
196
+
197
+ The steering hook is registered before generation and cleaned up after.
198
+ """
199
+ # Register hook if needed
200
+ if self.caa_enabled and self.steering_vector is not None and self._steering_hook_handle is None:
201
+ target_layer = self.model.layers[self.caa_layer_id]
202
+ self._steering_hook_handle = target_layer.register_forward_hook(self._apply_caa_steering)
203
+
204
+ try:
205
+ # Call parent generate method
206
+ outputs = super().generate(*args, **kwargs)
207
+ finally:
208
+ # Clean up hook after generation
209
+ if self._steering_hook_handle is not None:
210
+ self._steering_hook_handle.remove()
211
+ self._steering_hook_handle = None
212
+
213
+ return outputs
214
+
215
+ def set_caa_enabled(self, enabled: bool):
216
+ """Enable or disable CAA steering at runtime."""
217
+ self.caa_enabled = enabled
218
+ if not enabled and self._steering_hook_handle is not None:
219
+ self._steering_hook_handle.remove()
220
+ self._steering_hook_handle = None
221
+
222
+ def set_caa_alpha(self, alpha: float):
223
+ """Adjust CAA steering strength at runtime."""
224
+ self.caa_alpha = alpha
225
+
226
+ @classmethod
227
+ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
228
+ """
229
+ Load model with automatic CAA configuration.
230
+
231
+ This method ensures the steering vector is loaded from the embedded config.
232
+ If no weights are found locally, it loads from the base Qwen model.
233
+ """
234
+ import os
235
+ from pathlib import Path
236
+
237
+ # Check if we have local weights
238
+ local_path = Path(pretrained_model_name_or_path)
239
+ has_weights = any(
240
+ (local_path / f).exists()
241
+ for f in ["pytorch_model.bin", "model.safetensors", "pytorch_model.bin.index.json", "model.safetensors.index.json"]
242
+ )
243
+
244
+ if not has_weights and local_path.exists() and (local_path / "config.json").exists():
245
+ # We have config but no weights - load from base model
246
+ print(f"Loading weights from base model: Qwen/Qwen2.5-Coder-7B-Instruct")
247
+
248
+ # First, load config from local path
249
+ from transformers import AutoConfig
250
+ config = AutoConfig.from_pretrained(pretrained_model_name_or_path)
251
+
252
+ # Load model with base weights
253
+ # Remove config from kwargs if it exists to avoid conflict
254
+ kwargs_copy = kwargs.copy()
255
+ kwargs_copy.pop('config', None)
256
+
257
+ model = super().from_pretrained(
258
+ "Qwen/Qwen2.5-Coder-7B-Instruct",
259
+ *model_args,
260
+ config=config, # Use our custom config
261
+ **kwargs_copy
262
+ )
263
+
264
+ # Initialize CAA components
265
+ model.caa_enabled = config.caa_enabled
266
+ model.caa_layer_id = config.caa_layer_id
267
+ model.caa_alpha = config.caa_alpha
268
+ model.steering_method = config.steering_method
269
+ model._steering_hook_handle = None
270
+
271
+ # Load steering vector from config
272
+ if model.caa_enabled:
273
+ vector_path = config.steering_vector_path
274
+ if not os.path.isabs(vector_path):
275
+ vector_path = os.path.join(pretrained_model_name_or_path, vector_path)
276
+ model._load_steering_vector_from_file(vector_path)
277
+ else:
278
+ # Standard loading path
279
+ model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
280
+
281
+ # Load steering vector from config if not already loaded
282
+ if model.caa_enabled and model.steering_vector is None:
283
+ vector_path = model.config.steering_vector_path
284
+ if not os.path.isabs(vector_path):
285
+ vector_path = os.path.join(pretrained_model_name_or_path, vector_path)
286
+ model._load_steering_vector_from_file(vector_path)
287
+
288
+ return model
289
+
290
+
291
+ # Register the model
292
+ from transformers import AutoModelForCausalLM, AutoConfig
293
+
294
+ AutoConfig.register("wisent_qwen2", WisentQwen2Config)
295
+ AutoModelForCausalLM.register(WisentQwen2Config, WisentQwen2ForCausalLM)
special_tokens_map.json ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "additional_special_tokens": [
3
+ "<|im_start|>",
4
+ "<|im_end|>",
5
+ "<|object_ref_start|>",
6
+ "<|object_ref_end|>",
7
+ "<|box_start|>",
8
+ "<|box_end|>",
9
+ "<|quad_start|>",
10
+ "<|quad_end|>",
11
+ "<|vision_start|>",
12
+ "<|vision_end|>",
13
+ "<|vision_pad|>",
14
+ "<|image_pad|>",
15
+ "<|video_pad|>"
16
+ ],
17
+ "eos_token": {
18
+ "content": "<|im_end|>",
19
+ "lstrip": false,
20
+ "normalized": false,
21
+ "rstrip": false,
22
+ "single_word": false
23
+ },
24
+ "pad_token": {
25
+ "content": "<|endoftext|>",
26
+ "lstrip": false,
27
+ "normalized": false,
28
+ "rstrip": false,
29
+ "single_word": false
30
+ }
31
+ }
tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9c5ae00e602b8860cbd784ba82a8aa14e8feecec692e7076590d014d7b7fdafa
3
+ size 11421896
tokenizer_config.json ADDED
@@ -0,0 +1,207 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": false,
3
+ "add_prefix_space": false,
4
+ "added_tokens_decoder": {
5
+ "151643": {
6
+ "content": "<|endoftext|>",
7
+ "lstrip": false,
8
+ "normalized": false,
9
+ "rstrip": false,
10
+ "single_word": false,
11
+ "special": true
12
+ },
13
+ "151644": {
14
+ "content": "<|im_start|>",
15
+ "lstrip": false,
16
+ "normalized": false,
17
+ "rstrip": false,
18
+ "single_word": false,
19
+ "special": true
20
+ },
21
+ "151645": {
22
+ "content": "<|im_end|>",
23
+ "lstrip": false,
24
+ "normalized": false,
25
+ "rstrip": false,
26
+ "single_word": false,
27
+ "special": true
28
+ },
29
+ "151646": {
30
+ "content": "<|object_ref_start|>",
31
+ "lstrip": false,
32
+ "normalized": false,
33
+ "rstrip": false,
34
+ "single_word": false,
35
+ "special": true
36
+ },
37
+ "151647": {
38
+ "content": "<|object_ref_end|>",
39
+ "lstrip": false,
40
+ "normalized": false,
41
+ "rstrip": false,
42
+ "single_word": false,
43
+ "special": true
44
+ },
45
+ "151648": {
46
+ "content": "<|box_start|>",
47
+ "lstrip": false,
48
+ "normalized": false,
49
+ "rstrip": false,
50
+ "single_word": false,
51
+ "special": true
52
+ },
53
+ "151649": {
54
+ "content": "<|box_end|>",
55
+ "lstrip": false,
56
+ "normalized": false,
57
+ "rstrip": false,
58
+ "single_word": false,
59
+ "special": true
60
+ },
61
+ "151650": {
62
+ "content": "<|quad_start|>",
63
+ "lstrip": false,
64
+ "normalized": false,
65
+ "rstrip": false,
66
+ "single_word": false,
67
+ "special": true
68
+ },
69
+ "151651": {
70
+ "content": "<|quad_end|>",
71
+ "lstrip": false,
72
+ "normalized": false,
73
+ "rstrip": false,
74
+ "single_word": false,
75
+ "special": true
76
+ },
77
+ "151652": {
78
+ "content": "<|vision_start|>",
79
+ "lstrip": false,
80
+ "normalized": false,
81
+ "rstrip": false,
82
+ "single_word": false,
83
+ "special": true
84
+ },
85
+ "151653": {
86
+ "content": "<|vision_end|>",
87
+ "lstrip": false,
88
+ "normalized": false,
89
+ "rstrip": false,
90
+ "single_word": false,
91
+ "special": true
92
+ },
93
+ "151654": {
94
+ "content": "<|vision_pad|>",
95
+ "lstrip": false,
96
+ "normalized": false,
97
+ "rstrip": false,
98
+ "single_word": false,
99
+ "special": true
100
+ },
101
+ "151655": {
102
+ "content": "<|image_pad|>",
103
+ "lstrip": false,
104
+ "normalized": false,
105
+ "rstrip": false,
106
+ "single_word": false,
107
+ "special": true
108
+ },
109
+ "151656": {
110
+ "content": "<|video_pad|>",
111
+ "lstrip": false,
112
+ "normalized": false,
113
+ "rstrip": false,
114
+ "single_word": false,
115
+ "special": true
116
+ },
117
+ "151657": {
118
+ "content": "<tool_call>",
119
+ "lstrip": false,
120
+ "normalized": false,
121
+ "rstrip": false,
122
+ "single_word": false,
123
+ "special": false
124
+ },
125
+ "151658": {
126
+ "content": "</tool_call>",
127
+ "lstrip": false,
128
+ "normalized": false,
129
+ "rstrip": false,
130
+ "single_word": false,
131
+ "special": false
132
+ },
133
+ "151659": {
134
+ "content": "<|fim_prefix|>",
135
+ "lstrip": false,
136
+ "normalized": false,
137
+ "rstrip": false,
138
+ "single_word": false,
139
+ "special": false
140
+ },
141
+ "151660": {
142
+ "content": "<|fim_middle|>",
143
+ "lstrip": false,
144
+ "normalized": false,
145
+ "rstrip": false,
146
+ "single_word": false,
147
+ "special": false
148
+ },
149
+ "151661": {
150
+ "content": "<|fim_suffix|>",
151
+ "lstrip": false,
152
+ "normalized": false,
153
+ "rstrip": false,
154
+ "single_word": false,
155
+ "special": false
156
+ },
157
+ "151662": {
158
+ "content": "<|fim_pad|>",
159
+ "lstrip": false,
160
+ "normalized": false,
161
+ "rstrip": false,
162
+ "single_word": false,
163
+ "special": false
164
+ },
165
+ "151663": {
166
+ "content": "<|repo_name|>",
167
+ "lstrip": false,
168
+ "normalized": false,
169
+ "rstrip": false,
170
+ "single_word": false,
171
+ "special": false
172
+ },
173
+ "151664": {
174
+ "content": "<|file_sep|>",
175
+ "lstrip": false,
176
+ "normalized": false,
177
+ "rstrip": false,
178
+ "single_word": false,
179
+ "special": false
180
+ }
181
+ },
182
+ "additional_special_tokens": [
183
+ "<|im_start|>",
184
+ "<|im_end|>",
185
+ "<|object_ref_start|>",
186
+ "<|object_ref_end|>",
187
+ "<|box_start|>",
188
+ "<|box_end|>",
189
+ "<|quad_start|>",
190
+ "<|quad_end|>",
191
+ "<|vision_start|>",
192
+ "<|vision_end|>",
193
+ "<|vision_pad|>",
194
+ "<|image_pad|>",
195
+ "<|video_pad|>"
196
+ ],
197
+ "bos_token": null,
198
+ "clean_up_tokenization_spaces": false,
199
+ "eos_token": "<|im_end|>",
200
+ "errors": "replace",
201
+ "extra_special_tokens": {},
202
+ "model_max_length": 32768,
203
+ "pad_token": "<|endoftext|>",
204
+ "split_special_tokens": false,
205
+ "tokenizer_class": "Qwen2Tokenizer",
206
+ "unk_token": null
207
+ }
vectors/coding/steering_vector.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:35be276639827e6369810b3dcbfd30cdc61e8bf36abe77d61d9b7e904cc21088
3
+ size 7256
vocab.json ADDED
The diff for this file is too large to render. See raw diff