mirror of
https://github.com/PeiJieSun/DBDI-main.git
synced 2026-08-30 17:06:00 -07:00
PeijieSun
739 lines
30 KiB
Python
739 lines
30 KiB
Python
import torch
|
|
import json
|
|
import numpy as np
|
|
from typing import List, Dict, Tuple, Optional
|
|
from tqdm import tqdm
|
|
import argparse
|
|
from scipy import stats
|
|
import os
|
|
from sklearn.linear_model import LogisticRegression
|
|
from sklearn.model_selection import cross_val_score
|
|
from sklearn.preprocessing import StandardScaler
|
|
import matplotlib.pyplot as plt
|
|
import seaborn as sns
|
|
import gc
|
|
from datasets import load_dataset
|
|
import pickle
|
|
import time
|
|
from datetime import datetime
|
|
import pandas as pd
|
|
|
|
import sys
|
|
sys.path.append('.')
|
|
|
|
from config import model_paths
|
|
from utils import (
|
|
load_model, get_hidden_states, get_sentence_embeddings,
|
|
get_difference_matrix, get_svd, cosine_similarity
|
|
)
|
|
|
|
|
|
class ToxicVectorExtractor:
|
|
|
|
def __init__(self, model_name: str, device: str = 'cuda'):
|
|
self.device = device
|
|
self.model_name = model_name
|
|
|
|
print(f"Loading model: {model_name}")
|
|
self.model, self.tokenizer = load_model(model_name, model_paths)
|
|
|
|
self.num_layers = self.model.config.num_hidden_layers + 1
|
|
print(f"Model has {self.num_layers} layers")
|
|
|
|
def load_harmful_dataset(self, dataset_name: str, max_samples: int = None) -> List[str]:
|
|
"""
|
|
Loads a harmful dataset from a local JSON file.
|
|
|
|
Args:
|
|
dataset_name: Name of the dataset ('advbench', 'harmbench', 'strongreject').
|
|
max_samples: The maximum number of samples to use.
|
|
"""
|
|
print(f"\nLoading {dataset_name} dataset from local file...")
|
|
harmful_prompts = []
|
|
|
|
dataset_files = {
|
|
'advbench': './data/advbench.json',
|
|
'harmbench': './data/harmbench_validation.json',
|
|
'strongreject': './data/strongreject.json'
|
|
}
|
|
|
|
if dataset_name not in dataset_files:
|
|
print(f"Error: Unknown dataset {dataset_name}")
|
|
return []
|
|
|
|
filepath = dataset_files[dataset_name]
|
|
|
|
try:
|
|
if not os.path.exists(filepath):
|
|
print(f"Error: File {filepath} not found")
|
|
return []
|
|
|
|
with open(filepath, 'r', encoding='utf-8') as f:
|
|
data = json.load(f)
|
|
|
|
if isinstance(data, list):
|
|
for item in data:
|
|
if isinstance(item, dict) and 'mal' in item:
|
|
harmful_prompts.append(item['mal'])
|
|
elif isinstance(data, dict) and 'mal' in data:
|
|
harmful_prompts.append(data['mal'])
|
|
|
|
if max_samples is not None and harmful_prompts:
|
|
harmful_prompts = harmful_prompts[:max_samples]
|
|
|
|
print(f"Loaded {len(harmful_prompts)} harmful prompts from {dataset_name}")
|
|
|
|
if harmful_prompts:
|
|
print("Sample harmful prompts:")
|
|
for i, prompt in enumerate(harmful_prompts[:3]):
|
|
print(f" {i+1}: {prompt[:100]}...")
|
|
|
|
except Exception as e:
|
|
print(f"Error loading {dataset_name}: {e}")
|
|
harmful_prompts = []
|
|
|
|
return harmful_prompts
|
|
|
|
def load_harmless_dataset(self, max_samples: int = None) -> List[str]:
|
|
"""
|
|
Loads a harmless dataset from local files.
|
|
It prioritizes loading from './data/harmless.csv', then falls back to extracting
|
|
the 'benign' field from JSON files, and finally uses a default list.
|
|
"""
|
|
harmless_prompts = []
|
|
|
|
csv_filepath = './data/harmless.csv'
|
|
try:
|
|
if os.path.exists(csv_filepath):
|
|
print(f"Loading harmless prompts from {csv_filepath}")
|
|
df = pd.read_csv(csv_filepath)
|
|
|
|
if 'prompt' in df.columns:
|
|
harmless_prompts = df['prompt'].dropna().tolist()
|
|
print(f"Loaded {len(harmless_prompts)} harmless prompts from CSV")
|
|
else:
|
|
print(f"Warning: 'prompt' column not found in {csv_filepath}")
|
|
print(f"Available columns: {list(df.columns)}")
|
|
except Exception as e:
|
|
print(f"Error loading harmless.csv: {e}")
|
|
|
|
if not harmless_prompts:
|
|
json_filepath = './data/advbench.json'
|
|
try:
|
|
if os.path.exists(json_filepath):
|
|
print(f"Trying to load benign prompts from {json_filepath}")
|
|
with open(json_filepath, 'r', encoding='utf-8') as f:
|
|
data = json.load(f)
|
|
|
|
if isinstance(data, list):
|
|
for item in data:
|
|
if isinstance(item, dict) and 'benign' in item and item['benign']:
|
|
harmless_prompts.append(item['benign'])
|
|
|
|
if harmless_prompts:
|
|
print(f"Loaded {len(harmless_prompts)} benign prompts from JSON")
|
|
except Exception as e:
|
|
print(f"Error loading benign prompts from JSON: {e}")
|
|
|
|
if max_samples is not None:
|
|
harmless_prompts = harmless_prompts[:max_samples]
|
|
|
|
print(f"Final: Using {len(harmless_prompts)} harmless prompts")
|
|
|
|
return harmless_prompts
|
|
|
|
def collect_concept_activations(self, prompts: List[str], concept_name: str,
|
|
checkpoint_dir: Optional[str] = None) -> List[List[torch.Tensor]]:
|
|
"""
|
|
Collects activations for concept prompts, with checkpoint support.
|
|
"""
|
|
if not prompts:
|
|
print(f"Warning: No prompts provided for {concept_name}")
|
|
return [[] for _ in range(self.num_layers)]
|
|
|
|
if checkpoint_dir:
|
|
checkpoint_path = os.path.join(checkpoint_dir, f'{self.model_name}_{concept_name}_embeddings.pt')
|
|
if os.path.exists(checkpoint_path):
|
|
print(f"Loading {concept_name} embeddings from checkpoint...")
|
|
return torch.load(checkpoint_path)
|
|
|
|
print(f"\nCollecting {concept_name} concept activations for {len(prompts)} prompts...")
|
|
|
|
if len(prompts) > self.batch_size:
|
|
all_embeddings = [[] for _ in range(self.num_layers)]
|
|
|
|
for i in tqdm(range(0, len(prompts), self.batch_size), desc=f"Processing {concept_name} batches"):
|
|
batch_prompts = prompts[i:i+self.batch_size]
|
|
|
|
try:
|
|
batch_embeddings = get_sentence_embeddings(
|
|
batch_prompts, self.model, self.model_name, self.tokenizer
|
|
)
|
|
|
|
for layer_idx in range(self.num_layers):
|
|
all_embeddings[layer_idx].extend(batch_embeddings[layer_idx])
|
|
|
|
if i % (self.batch_size * 4) == 0:
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
except Exception as e:
|
|
print(f"Warning: Error processing batch {i//self.batch_size}: {e}")
|
|
continue
|
|
|
|
embeddings = all_embeddings
|
|
else:
|
|
embeddings = get_sentence_embeddings(
|
|
prompts, self.model, self.model_name, self.tokenizer
|
|
)
|
|
|
|
if checkpoint_dir:
|
|
os.makedirs(checkpoint_dir, exist_ok=True)
|
|
torch.save(embeddings, checkpoint_path)
|
|
print(f"Saved {concept_name} embeddings checkpoint")
|
|
|
|
return embeddings
|
|
|
|
def train_toxic_classifiers(self,
|
|
toxic_embeddings: List[List[torch.Tensor]],
|
|
neutral_embeddings: List[List[torch.Tensor]]) -> Dict[int, Dict]:
|
|
"""
|
|
Trains a classifier for each layer to evaluate its ability to distinguish toxic/neutral concepts.
|
|
"""
|
|
layer_performance = {}
|
|
|
|
print("\nTraining toxic concept classifiers for each layer...")
|
|
for layer_idx in tqdm(range(self.num_layers)):
|
|
try:
|
|
X_toxic = torch.stack(toxic_embeddings[layer_idx]).cpu().numpy()
|
|
X_neutral = torch.stack(neutral_embeddings[layer_idx]).cpu().numpy()
|
|
|
|
X = np.vstack([X_toxic, X_neutral])
|
|
y = np.hstack([np.ones(len(X_toxic)), np.zeros(len(X_neutral))])
|
|
|
|
scaler = StandardScaler()
|
|
X_scaled = scaler.fit_transform(X)
|
|
|
|
if 'deepseek' in self.model_name:
|
|
clf = LogisticRegression(
|
|
max_iter=2000,
|
|
random_state=42,
|
|
class_weight='balanced',
|
|
C=0.1
|
|
)
|
|
else:
|
|
clf = LogisticRegression(
|
|
max_iter=1000,
|
|
random_state=42,
|
|
class_weight='balanced'
|
|
)
|
|
|
|
scores = cross_val_score(clf, X_scaled, y, cv=5, scoring='accuracy')
|
|
|
|
clf.fit(X_scaled, y)
|
|
|
|
feature_importance = np.abs(clf.coef_[0])
|
|
|
|
t_stats = []
|
|
p_values = []
|
|
for i in range(X_toxic.shape[1]):
|
|
try:
|
|
t_stat, p_val = stats.ttest_ind(X_toxic[:, i], X_neutral[:, i])
|
|
t_stats.append(abs(t_stat))
|
|
p_values.append(p_val)
|
|
except:
|
|
t_stats.append(0)
|
|
p_values.append(1)
|
|
|
|
layer_performance[layer_idx] = {
|
|
'accuracy': scores.mean(),
|
|
'accuracy_std': scores.std(),
|
|
'classifier': clf,
|
|
'scaler': scaler,
|
|
'feature_importance': feature_importance,
|
|
't_stats': np.array(t_stats),
|
|
'p_values': np.array(p_values),
|
|
'top_features': np.argsort(feature_importance)[-min(100, len(feature_importance)):]
|
|
}
|
|
|
|
except Exception as e:
|
|
print(f"Warning: Failed to train classifier for layer {layer_idx}: {e}")
|
|
layer_performance[layer_idx] = {
|
|
'accuracy': 0.5,
|
|
'accuracy_std': 0.0,
|
|
'classifier': None,
|
|
'scaler': None,
|
|
'feature_importance': None,
|
|
't_stats': None,
|
|
'p_values': None,
|
|
'top_features': None
|
|
}
|
|
|
|
return layer_performance
|
|
|
|
def select_best_toxic_layers(self, layer_performance: Dict, top_k: int = 5) -> List[int]:
|
|
"""
|
|
Selects the most suitable layers for extracting the toxic concept.
|
|
"""
|
|
valid_layers = {k: v for k, v in layer_performance.items()
|
|
if v['accuracy'] > 0.5 and v['classifier'] is not None}
|
|
|
|
if not valid_layers:
|
|
print("Warning: No valid layers found, using default layers")
|
|
return list(range(max(0, self.num_layers - top_k), self.num_layers))
|
|
|
|
sorted_layers = sorted(
|
|
valid_layers.items(),
|
|
key=lambda x: x[1]['accuracy'],
|
|
reverse=True
|
|
)
|
|
|
|
print("\nLayer toxic classification performance:")
|
|
for layer_idx, perf in sorted_layers[:10]:
|
|
print(f"Layer {layer_idx}: Accuracy = {perf['accuracy']:.4f} (±{perf['accuracy_std']:.4f})")
|
|
|
|
best_layers = [layer_idx for layer_idx, _ in sorted_layers[:min(top_k, len(sorted_layers))]]
|
|
|
|
if len(best_layers) < top_k:
|
|
for layer in range(self.num_layers - 1, -1, -1):
|
|
if layer not in best_layers:
|
|
best_layers.append(layer)
|
|
if len(best_layers) >= top_k:
|
|
break
|
|
|
|
return best_layers[:top_k]
|
|
|
|
def extract_toxic_vector_for_layer(self,
|
|
toxic_embeddings: List[torch.Tensor],
|
|
neutral_embeddings: List[torch.Tensor],
|
|
layer_performance: Dict,
|
|
method: str = 'dbdi',
|
|
neuron_ratio: float = 0.25) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Extracts the toxic vector for a single layer.
|
|
"""
|
|
try:
|
|
if method == 'dbdi':
|
|
difference_matrix = get_difference_matrix(toxic_embeddings, neutral_embeddings)
|
|
_, S, V = get_svd(difference_matrix)
|
|
toxic_direction = V[:, 0].cpu()
|
|
elif method == 'mean_diff':
|
|
toxic_mean = torch.stack(toxic_embeddings).mean(dim=0)
|
|
neutral_mean = torch.stack(neutral_embeddings).mean(dim=0)
|
|
toxic_direction = (toxic_mean - neutral_mean).cpu()
|
|
|
|
if (layer_performance.get('p_values') is None or
|
|
layer_performance.get('t_stats') is None or
|
|
layer_performance.get('feature_importance') is None):
|
|
|
|
abs_direction = torch.abs(toxic_direction)
|
|
k = int(len(toxic_direction) * neuron_ratio)
|
|
top_k_indices = torch.topk(abs_direction, k).indices
|
|
combined_mask = torch.zeros_like(toxic_direction, dtype=torch.bool)
|
|
combined_mask[top_k_indices] = True
|
|
|
|
else:
|
|
p_values = layer_performance['p_values']
|
|
t_stats = layer_performance['t_stats']
|
|
feature_importance = torch.tensor(layer_performance['feature_importance'])
|
|
|
|
significance_mask = torch.tensor(p_values < 0.01)
|
|
|
|
t_stats_tensor = torch.tensor(t_stats)
|
|
if torch.sum(t_stats_tensor > 0) > 0:
|
|
effect_threshold = torch.quantile(t_stats_tensor[t_stats_tensor > 0], 1 - neuron_ratio)
|
|
effect_mask = t_stats_tensor > effect_threshold
|
|
else:
|
|
effect_mask = torch.zeros_like(significance_mask)
|
|
|
|
if torch.sum(feature_importance > 0) > 0:
|
|
importance_threshold = torch.quantile(feature_importance[feature_importance > 0], 1 - neuron_ratio)
|
|
importance_mask = feature_importance > importance_threshold
|
|
else:
|
|
importance_mask = torch.zeros_like(significance_mask)
|
|
|
|
combined_mask = significance_mask & (effect_mask | importance_mask)
|
|
|
|
if combined_mask.sum() < int(len(toxic_direction) * neuron_ratio):
|
|
k = int(len(toxic_direction) * neuron_ratio)
|
|
if torch.sum(t_stats_tensor > 0) > 0:
|
|
top_k_indices = torch.topk(t_stats_tensor, min(k, len(t_stats_tensor))).indices
|
|
else:
|
|
top_k_indices = torch.topk(torch.abs(toxic_direction), k).indices
|
|
combined_mask = torch.zeros_like(toxic_direction, dtype=torch.bool)
|
|
combined_mask[top_k_indices] = True
|
|
|
|
sparse_toxic_vector = toxic_direction * combined_mask.float()
|
|
|
|
if torch.norm(sparse_toxic_vector) > 1e-8:
|
|
sparse_toxic_vector = sparse_toxic_vector / torch.norm(sparse_toxic_vector)
|
|
|
|
return sparse_toxic_vector, combined_mask
|
|
|
|
except Exception as e:
|
|
print(f"Error in extract_toxic_vector_for_layer: {e}")
|
|
dim = len(toxic_embeddings[0])
|
|
return torch.zeros(dim), torch.zeros(dim, dtype=torch.bool)
|
|
|
|
def extract_vectors_for_dataset(self,
|
|
dataset_name: str,
|
|
max_samples: int = None,
|
|
extraction_method: str = 'dbdi',
|
|
neuron_ratio: float = 0.25,
|
|
checkpoint_dir: Optional[str] = None) -> Dict:
|
|
"""
|
|
Extracts vectors for a single dataset.
|
|
"""
|
|
harmful_prompts = self.load_harmful_dataset(dataset_name, max_samples)
|
|
if not harmful_prompts:
|
|
print(f"Skipping {dataset_name} - no data loaded")
|
|
return None
|
|
|
|
harmless_prompts = self.load_harmless_dataset(len(harmful_prompts))
|
|
|
|
harmful_embeddings = self.collect_concept_activations(
|
|
harmful_prompts, f"{dataset_name}_harmful", checkpoint_dir
|
|
)
|
|
harmless_embeddings = self.collect_concept_activations(
|
|
harmless_prompts, f"{dataset_name}_harmless", checkpoint_dir
|
|
)
|
|
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
layer_perf_checkpoint_path = None
|
|
if checkpoint_dir:
|
|
layer_perf_checkpoint_path = os.path.join(
|
|
checkpoint_dir,
|
|
f'{self.model_name}_{dataset_name}_layer_performance.pkl'
|
|
)
|
|
if os.path.exists(layer_perf_checkpoint_path):
|
|
print(f"Loading layer performance checkpoint for {dataset_name}...")
|
|
with open(layer_perf_checkpoint_path, 'rb') as f:
|
|
layer_performance = pickle.load(f)
|
|
else:
|
|
layer_performance = None
|
|
else:
|
|
layer_performance = None
|
|
|
|
if layer_performance is None:
|
|
layer_performance = self.train_toxic_classifiers(harmful_embeddings, harmless_embeddings)
|
|
|
|
if checkpoint_dir and layer_perf_checkpoint_path:
|
|
with open(layer_perf_checkpoint_path, 'wb') as f:
|
|
pickle.dump(layer_performance, f)
|
|
print(f"Saved layer performance checkpoint for {dataset_name}")
|
|
|
|
best_layers = self.select_best_toxic_layers(layer_performance, top_k=5)
|
|
|
|
results = {
|
|
'dataset_name': dataset_name,
|
|
'layer_performance': {},
|
|
'harmful_vectors': {},
|
|
'best_layers': best_layers,
|
|
'n_prompts': len(harmful_prompts)
|
|
}
|
|
|
|
for layer_idx, perf in layer_performance.items():
|
|
results['layer_performance'][layer_idx] = {
|
|
'accuracy': perf['accuracy'],
|
|
'accuracy_std': perf['accuracy_std'],
|
|
'n_significant_features': int((perf['p_values'] < 0.01).sum()) if perf['p_values'] is not None else 0,
|
|
'mean_t_stat': float(perf['t_stats'].mean()) if perf['t_stats'] is not None else 0.0
|
|
}
|
|
|
|
print(f"\nExtracting harmful vectors for all {self.num_layers} layers...")
|
|
for layer_idx in tqdm(range(self.num_layers)):
|
|
try:
|
|
vector, mask = self.extract_toxic_vector_for_layer(
|
|
harmful_embeddings[layer_idx],
|
|
harmless_embeddings[layer_idx],
|
|
layer_performance[layer_idx],
|
|
method=extraction_method,
|
|
neuron_ratio=neuron_ratio
|
|
)
|
|
|
|
results['harmful_vectors'][layer_idx] = {
|
|
'vector': vector.cpu(),
|
|
'mask': mask.cpu(),
|
|
'n_active': mask.sum().item(),
|
|
'sparsity': 1 - (mask.sum().item() / len(mask)),
|
|
'is_best_layer': layer_idx in best_layers
|
|
}
|
|
except Exception as e:
|
|
print(f"Warning: Failed to extract vector for layer {layer_idx}: {e}")
|
|
|
|
return results
|
|
|
|
def save_dataset_results(self, results: Dict, output_dir: str, dataset_name: str):
|
|
"""
|
|
Saves the results for a single dataset.
|
|
"""
|
|
dataset_output_dir = os.path.join(output_dir, dataset_name)
|
|
os.makedirs(dataset_output_dir, exist_ok=True)
|
|
|
|
layer_accuracy_path = os.path.join(dataset_output_dir, 'layer_accuracies.json')
|
|
layer_accuracies = {}
|
|
for layer_idx, perf in results['layer_performance'].items():
|
|
layer_accuracies[f"layer_{layer_idx}"] = {
|
|
'accuracy': perf['accuracy'],
|
|
'accuracy_std': perf['accuracy_std'],
|
|
'n_significant_features': perf['n_significant_features'],
|
|
'mean_t_stat': perf['mean_t_stat']
|
|
}
|
|
|
|
with open(layer_accuracy_path, 'w') as f:
|
|
json.dump(layer_accuracies, f, indent=2)
|
|
|
|
metadata = {
|
|
'model_name': self.model_name,
|
|
'dataset_name': dataset_name,
|
|
'num_layers': self.num_layers,
|
|
'best_layers': results['best_layers'],
|
|
'n_prompts': results['n_prompts'],
|
|
'extraction_time': datetime.now().isoformat()
|
|
}
|
|
|
|
with open(os.path.join(dataset_output_dir, 'metadata.json'), 'w') as f:
|
|
json.dump(metadata, f, indent=2)
|
|
|
|
all_vectors_data = {
|
|
'harmful_vectors': results['harmful_vectors'],
|
|
'metadata': metadata
|
|
}
|
|
|
|
torch.save(all_vectors_data, os.path.join(dataset_output_dir, 'all_layer_vectors.pt'))
|
|
|
|
best_vectors = {}
|
|
for layer_idx in results['best_layers']:
|
|
if layer_idx in results['harmful_vectors']:
|
|
best_vectors[layer_idx] = results['harmful_vectors'][layer_idx]
|
|
|
|
torch.save(best_vectors, os.path.join(dataset_output_dir, 'best_5_layer_vectors.pt'))
|
|
|
|
self.plot_layer_performance(results, dataset_output_dir)
|
|
|
|
print(f"\nResults for {dataset_name} saved to {dataset_output_dir}")
|
|
print(f"- layer_accuracies.json: Per-layer accuracy metrics")
|
|
print(f"- metadata.json: Dataset and extraction metadata")
|
|
print(f"- all_layer_vectors.pt: Vectors from all layers")
|
|
print(f"- best_5_layer_vectors.pt: Vectors from top 5 layers")
|
|
|
|
def plot_layer_performance(self, results: Dict, output_dir: str):
|
|
"""
|
|
Creates a visualization of layer performance.
|
|
"""
|
|
layers = sorted(results['layer_performance'].keys())
|
|
accuracies = [results['layer_performance'][l]['accuracy'] for l in layers]
|
|
stds = [results['layer_performance'][l]['accuracy_std'] for l in layers]
|
|
|
|
plt.figure(figsize=(12, 6))
|
|
plt.errorbar(layers, accuracies, yerr=stds, marker='o', capsize=5, color='red')
|
|
plt.xlabel('Layer Index')
|
|
plt.ylabel('Classification Accuracy')
|
|
plt.title(f'Layer-wise Harmful Classification Performance\nModel: {self.model_name}, Dataset: {results["dataset_name"]}')
|
|
plt.grid(True, alpha=0.3)
|
|
|
|
for layer in results['best_layers']:
|
|
plt.axvline(x=layer, color='darkred', linestyle='--', alpha=0.5)
|
|
|
|
plt.axhline(y=0.5, color='gray', linestyle=':', alpha=0.5, label='Random baseline')
|
|
plt.legend()
|
|
plt.tight_layout()
|
|
plt.savefig(os.path.join(output_dir, 'layer_performance.png'))
|
|
plt.close()
|
|
|
|
def cleanup(self):
|
|
"""Cleans up the model to free up memory."""
|
|
if hasattr(self, 'model'):
|
|
del self.model
|
|
if hasattr(self, 'tokenizer'):
|
|
del self.tokenizer
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
|
|
class MultiModelMultiDatasetExtractor:
|
|
"""Extractor to handle multiple models and multiple datasets."""
|
|
|
|
def __init__(self, models: List[str], datasets: List[str], device: str = 'cuda'):
|
|
self.models = models
|
|
self.datasets = datasets
|
|
self.device = device
|
|
self.checkpoint_dir = './checkpoints_harm'
|
|
os.makedirs(self.checkpoint_dir, exist_ok=True)
|
|
|
|
def get_progress_checkpoint_path(self):
|
|
"""Gets the path for the progress checkpoint file."""
|
|
return os.path.join(self.checkpoint_dir, 'extraction_progress.json')
|
|
|
|
def load_progress(self) -> Dict:
|
|
"""Loads the extraction progress."""
|
|
checkpoint_path = self.get_progress_checkpoint_path()
|
|
if os.path.exists(checkpoint_path):
|
|
with open(checkpoint_path, 'r') as f:
|
|
return json.load(f)
|
|
return {'completed_tasks': [], 'timestamp': datetime.now().isoformat()}
|
|
|
|
def save_progress(self, progress: Dict):
|
|
"""Saves the extraction progress."""
|
|
checkpoint_path = self.get_progress_checkpoint_path()
|
|
progress['timestamp'] = datetime.now().isoformat()
|
|
with open(checkpoint_path, 'w') as f:
|
|
json.dump(progress, f, indent=2)
|
|
|
|
def process_all(self, max_samples: int = None, neuron_ratio: float = 0.25):
|
|
"""
|
|
Processes all combinations of models and datasets.
|
|
"""
|
|
progress = self.load_progress()
|
|
completed_tasks = progress.get('completed_tasks', [])
|
|
|
|
total_tasks = len(self.models) * len(self.datasets)
|
|
print(f"\n{'='*60}")
|
|
print(f"Starting multi-model multi-dataset extraction")
|
|
print(f"Models: {self.models}")
|
|
print(f"Datasets: {self.datasets}")
|
|
print(f"Total tasks: {total_tasks}")
|
|
print(f"Completed tasks: {len(completed_tasks)}")
|
|
print(f"{'='*60}\n")
|
|
|
|
start_time = time.time()
|
|
task_idx = 0
|
|
|
|
for model_name in self.models:
|
|
try:
|
|
extractor = ToxicVectorExtractor(model_name, self.device)
|
|
except Exception as e:
|
|
print(f"Failed to load model {model_name}: {e}")
|
|
continue
|
|
|
|
output_dir = os.path.join('./extracted_harm_vector', model_name)
|
|
|
|
for dataset_name in self.datasets:
|
|
task_idx += 1
|
|
task_key = f"{model_name}_{dataset_name}"
|
|
|
|
if task_key in completed_tasks:
|
|
print(f"\n[{task_idx}/{total_tasks}] Task {task_key} already completed, skipping...")
|
|
continue
|
|
|
|
print(f"\n{'='*60}")
|
|
print(f"[{task_idx}/{total_tasks}] Processing: {model_name} + {dataset_name}")
|
|
print(f"{'='*60}")
|
|
|
|
try:
|
|
results = extractor.extract_vectors_for_dataset(
|
|
dataset_name=dataset_name,
|
|
max_samples=max_samples,
|
|
extraction_method='dbdi',
|
|
neuron_ratio=neuron_ratio,
|
|
checkpoint_dir=self.checkpoint_dir
|
|
)
|
|
|
|
if results:
|
|
extractor.save_dataset_results(results, output_dir, dataset_name)
|
|
|
|
completed_tasks.append(task_key)
|
|
progress['completed_tasks'] = completed_tasks
|
|
self.save_progress(progress)
|
|
|
|
print(f"\n✓ Task {task_key} completed successfully")
|
|
else:
|
|
print(f"\n✗ Task {task_key} skipped - no data")
|
|
|
|
except Exception as e:
|
|
print(f"\n✗ Error in task {task_key}: {str(e)}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
if task_idx < total_tasks:
|
|
response = input("\nContinue with next task? (y/n): ")
|
|
if response.lower() != 'y':
|
|
break
|
|
|
|
extractor.cleanup()
|
|
print(f"\nCompleted all datasets for model {model_name}")
|
|
|
|
elapsed_time = time.time() - start_time
|
|
print(f"\n{'='*60}")
|
|
print(f"Multi-model multi-dataset extraction completed!")
|
|
print(f"Total time: {elapsed_time/60:.2f} minutes")
|
|
print(f"Completed tasks: {len(completed_tasks)}/{total_tasks}")
|
|
print(f"{'='*60}")
|
|
|
|
if len(completed_tasks) == total_tasks:
|
|
response = input("\nAll tasks completed. Remove checkpoints? (y/n): ")
|
|
if response.lower() == 'y':
|
|
self.cleanup_checkpoints()
|
|
|
|
def cleanup_checkpoints(self):
|
|
"""Cleans up all checkpoint files."""
|
|
print("\nCleaning up checkpoints...")
|
|
import shutil
|
|
if os.path.exists(self.checkpoint_dir):
|
|
shutil.rmtree(self.checkpoint_dir)
|
|
print(f"Removed checkpoint directory: {self.checkpoint_dir}")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description='Extract harmful concept vectors using multiple datasets')
|
|
parser.add_argument('--models', type=str, nargs='+',
|
|
choices=list(model_paths.keys()),
|
|
help='Model names (can specify multiple)')
|
|
parser.add_argument('--model', type=str,
|
|
choices=list(model_paths.keys()),
|
|
help='Single model name (for backward compatibility)')
|
|
parser.add_argument('--datasets', type=str, nargs='+',
|
|
default=['advbench', 'harmbench', 'strongreject'],
|
|
choices=['advbench', 'harmbench', 'strongreject'],
|
|
help='Datasets to use for harmful prompts')
|
|
parser.add_argument('--max_samples', type=int, default=None,
|
|
help='Maximum samples from each dataset')
|
|
parser.add_argument('--neuron_ratio', type=float, default=0.25,
|
|
help='Ratio of neurons to select (default: 0.25)')
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.models:
|
|
models_to_process = args.models
|
|
elif args.model:
|
|
models_to_process = [args.model]
|
|
else:
|
|
models_to_process = list(model_paths.keys())
|
|
print(f"No models specified, will process all available models: {models_to_process}")
|
|
response = input("Continue? (y/n): ")
|
|
if response.lower() != 'y':
|
|
return
|
|
|
|
if len(models_to_process) == 1 and len(args.datasets) == 1:
|
|
model_name = models_to_process[0]
|
|
dataset_name = args.datasets[0]
|
|
|
|
extractor = ToxicVectorExtractor(model_name)
|
|
output_dir = os.path.join('./extracted_harm_vector', model_name)
|
|
|
|
print(f"\nExtracting harmful vectors for {model_name} using {dataset_name}...")
|
|
|
|
results = extractor.extract_vectors_for_dataset(
|
|
dataset_name=dataset_name,
|
|
max_samples=args.max_samples,
|
|
neuron_ratio=args.neuron_ratio
|
|
)
|
|
|
|
if results:
|
|
extractor.save_dataset_results(results, output_dir, dataset_name)
|
|
|
|
else:
|
|
multi_extractor = MultiModelMultiDatasetExtractor(
|
|
models=models_to_process,
|
|
datasets=args.datasets
|
|
)
|
|
multi_extractor.process_all(
|
|
max_samples=args.max_samples,
|
|
neuron_ratio=args.neuron_ratio
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |