Parameter Optimization

direction_scope = trial.suggest_categorical(
    "direction_scope",
    ["global", "per layer"],
)

direction_index = trial.suggest_float(
    "direction_index",
    0.4 * last_layer_index,  # Search from 40% through layers
    0.9 * last_layer_index,  # Up to 90% through layers
)

max_weight = trial.suggest_float(
    f"{component}.max_weight",
    0.8,   # Minimum ablation strength
    1.5,   # Maximum ablation strength
)

max_weight_position = trial.suggest_float(
    f"{component}.max_weight_position",
    0.6 * last_layer_index,  # Later layers
    1.0 * last_layer_index,  # Up to final layer
)

# Sampled as fraction of max_weight for multivariate TPE
min_weight_fraction = trial.suggest_float(
    f"{component}.min_weight",
    0.0,   # Complete taper to zero
    1.0,   # Constant weight across all layers
)
min_weight = min_weight_fraction * max_weight

min_weight_distance = trial.suggest_float(
    f"{component}.min_weight_distance",
    1.0,                      # Single layer only
    0.6 * last_layer_index,   # Up to 60% of all layers
)

for layer_index in range(len(self.get_layers())):
    distance = abs(layer_index - params.max_weight_position)
    
    # Don't orthogonalize layers outside the kernel
    if distance > params.min_weight_distance:
        continue
    
    # Linear interpolation from max_weight to min_weight
    weight = params.max_weight + (
        distance / params.min_weight_distance
    ) * (params.min_weight - params.max_weight)

logprobs = self.model.get_logprobs_batched(self.good_prompts)
kl_divergence = F.kl_div(
    logprobs,           # Abliterated model
    self.base_logprobs, # Original model
    reduction="batchmean",
    log_target=True,
).item()

def is_refusal(self, response: str) -> bool:
    if not response.strip():
        return True  # Empty responses are refusals
    
    # Normalize text
    response = response.lower().replace("*", "")
    response = response.replace("'", "'")  # Typographic apostrophe
    response = " ".join(response.split())  # Normalize whitespace
    
    # Check for refusal markers
    for marker in self.settings.refusal_markers:
        if marker.lower() in response:
            return True
    
    return False

# Get Pareto-optimal trials
sorted_trials = sorted(
    completed_trials,
    key=lambda trial: (
        trial.user_attrs["refusals"],
        trial.user_attrs["kl_divergence"],
    ),
)
min_divergence = math.inf
best_trials = []
for trial in sorted_trials:
    kl_divergence = trial.user_attrs["kl_divergence"]
    if kl_divergence < min_divergence:
        min_divergence = kl_divergence
        best_trials.append(trial)

study_checkpoint_file = os.path.join(
    settings.study_checkpoint_dir,
    "".join([
        (c if (c.isalnum() or c in ["_", "-"]) else "--") 
        for c in settings.model
    ]) + ".jsonl",
)

backend = JournalFileBackend(study_checkpoint_file)
storage = JournalStorage(backend)

if settings.batch_size == 0:
    print("Determining optimal batch size...")
    
    batch_size = 1
    best_batch_size = -1
    best_performance = -1
    
    while batch_size <= settings.max_batch_size:
        try:
            # Warmup run
            model.get_responses(prompts)
            
            # Benchmark run
            start_time = time.perf_counter()
            responses = model.get_responses(prompts)
            end_time = time.perf_counter()
            
            performance = sum(response_lengths) / (end_time - start_time)
            
            if performance > best_performance:
                best_batch_size = batch_size
                best_performance = performance
            
            batch_size *= 2
        except Exception:
            break
    
    settings.batch_size = best_batch_size

quantization = "bnb_4bit"