Use black to format

hacking.py

@@ -10,7 +10,16 @@ from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassifica
 from huggingface_hub import login
 from llm_attacks.minimal_gcg.opt_utils import get_filtered_cands
 from words import words
-from utils import minimize_tokens, sample_control, count_tokens, get_combined_score, get_random_words, token_gradients_combined, find_best_word_to_add, words_db
+from utils import (
+    minimize_tokens,
+    sample_control,
+    count_tokens,
+    get_combined_score,
+    get_random_words,
+    token_gradients_combined,
+    find_best_word_to_add,
+    words_db,
+)
 
 # check if cuda is available
 cuda_available: bool = torch.cuda.is_available()
@@ -19,7 +28,7 @@ device: torch.device = torch.device("cuda" if cuda_available else "cpu")
 # use token from environment variable
 login(token=os.getenv("HF_TOKEN"))
 
-alpha: float = 0.3 # weight between benign maximization and loss minimization
+alpha: float = 0.3  # weight between benign maximization and loss minimization
 min_benign_confidence: float = 0.95  # Minimum benign probability to stop (95%)
 words_to_inject: int = 1
 improvement_threshold: float = 0.999
@@ -27,7 +36,9 @@ rolling_window_size: int = 10  # Number of past scores to consider
 
 # Define a threshold for injecting educational text when optimization stagnates
 stagnation_threshold: int = 15  # Inject after this many iterations without significant improvement
-grace_period: int = 50  # Number of iterations to allow after text injection before counting against patience
+grace_period: int = (
+    50  # Number of iterations to allow after text injection before counting against patience
+)
 max_iterations: int = 10000  # Maximum number of optimization iterations
 patience: int = 10000  # Number of iterations with no improvement before early stopping
 max_top_scores: int = 10  # Number of top scores to maintain
@@ -39,9 +50,11 @@ token_penalty_weight: float = 0.4
 min_acceptable_benign: float = 0.50
 
 model_id: str = "meta-llama/Llama-Prompt-Guard-2-86M"
-#model_id = "meta-llama/Llama-Prompt-Guard-2-22M"
+# model_id = "meta-llama/Llama-Prompt-Guard-2-22M"
 tokenizer: AutoTokenizer = AutoTokenizer.from_pretrained(model_id)
-model: AutoModelForSequenceClassification = AutoModelForSequenceClassification.from_pretrained(model_id)
+model: AutoModelForSequenceClassification = AutoModelForSequenceClassification.from_pretrained(
+    model_id
+)
 model = model.to(device)  # Move model to MPS device
 
 benign_class: str = "label_0"
@@ -72,37 +85,51 @@ init_token_priority: float = 0.0
 general_token_priority: float = 0.95
 
 # Try to use top-performing words from the database for the initial prefix
-top_words = words_db.get_top_words(limit=init_prefix_words_count, min_uses=1, token_weight=init_token_priority)
+top_words = words_db.get_top_words(
+    limit=init_prefix_words_count, min_uses=1, token_weight=init_token_priority
+)
 if top_words:
     print(f"Using {len(top_words)} top-performing words from database for initial prefix")
     # Get words with combined token and improvement prioritization
     initial_words = get_random_words(
         n=init_prefix_words_count,
         min_uses=1,  # Words must have been tested at least once
-        token_priority=init_token_priority
+        token_priority=init_token_priority,
     )
     init_prefix: str = " ".join(initial_words)
-    print(f"Created initial prefix using database-informed words (token priority: {init_token_priority})")
+    print(
+        f"Created initial prefix using database-informed words (token priority: {init_token_priority})"
+    )
 else:
     # Fall back to random words if the database doesn't have enough data
     init_prefix: str = " ".join(words[:init_prefix_words_count])
     print(f"Using random words for initial prefix (no database history available)")
 
-#init_prefix = "".join(random.choices(words, k=init_prefix_words_count))
+# init_prefix = "".join(random.choices(words, k=init_prefix_words_count))
+
 
 def main():
     global injection_text, text, init_prefix_words_count
     # Parse command line arguments
     parser = argparse.ArgumentParser(description="Prompt hacking tool")
-    parser.add_argument("--injection", type=str, 
+    parser.add_argument(
-                        default=injection_text,
+        "--injection",
-                        help="Injection text to use in the template")
+        type=str,
-    parser.add_argument("--mandatory-text", type=str,
+        default=injection_text,
-                        default=text,
+        help="Injection text to use in the template",
-                        help="Mandatory text to use in the template")
+    )
-    parser.add_argument("--init-prefix-words-count", type=int,
+    parser.add_argument(
-                        default=init_prefix_words_count,
+        "--mandatory-text",
-                        help="Number of words to use in the initial prefix")
+        type=str,
+        default=text,
+        help="Mandatory text to use in the template",
+    )
+    parser.add_argument(
+        "--init-prefix-words-count",
+        type=int,
+        default=init_prefix_words_count,
+        help="Number of words to use in the initial prefix",
+    )
 
     args = parser.parse_args()
 
@@ -117,14 +144,16 @@ def main():
     print(f"\nTrying initial prefix: {init_prefix}")
 
     # Convert initial adversarial string to tokens
-    best_score: float = float('-inf')
+    best_score: float = float("-inf")
     best_prefix: Optional[str] = None
     adv_prefix: str = init_prefix
-    adv_prefix_tokens: torch.Tensor = tokenizer(adv_prefix, return_tensors="pt", add_special_tokens=False)["input_ids"][0]
+    adv_prefix_tokens: torch.Tensor = tokenizer(
+        adv_prefix, return_tensors="pt", add_special_tokens=False
+    )["input_ids"][0]
     adv_prefix_tokens = adv_prefix_tokens.to(device)  # Move tokens to MPS device
     control_slice: slice = slice(0, len(adv_prefix_tokens))  # Slice representing the prefix tokens
 
-    best_iteration_score: float = float('-inf')
+    best_iteration_score: float = float("-inf")
     iterations_without_improvement: int = 0
 
     # Track both rolling and top scores
@@ -139,7 +168,7 @@ def main():
         # Prepare input tensors using template
         full_text = injection_text + adv_prefix + text
         inputs: Dict[str, torch.Tensor] = tokenizer(full_text, return_tensors="pt")
-        input_ids: torch.Tensor = inputs['input_ids'][0].to(device)  # Move input_ids to MPS device
+        input_ids: torch.Tensor = inputs["input_ids"][0].to(device)  # Move input_ids to MPS device
 
         # Compute gradients using combined approach
         coordinate_grad: torch.Tensor = token_gradients_combined(
@@ -149,7 +178,7 @@ def main():
             benign_class=benign_class_idx,
             malicious_class=malicious_class_idx,
             alpha=alpha,
-            device=device
+            device=device,
         )
 
         # Ensure coordinate_grad is on the correct device and has the right shape
@@ -164,9 +193,9 @@ def main():
             new_adv_prefix_toks: torch.Tensor = sample_control(
                 adv_prefix_tokens,
                 coordinate_grad,
-                batch_size=32,      # Larger batch for more candidates
+                batch_size=32,  # Larger batch for more candidates
-                topk=16,            # More options per token
+                topk=16,  # More options per token
-                temp=1.5,            # Higher temperature for more exploration
+                temp=1.5,  # Higher temperature for more exploration
             )
 
             # Convert new tokens to text
@@ -174,7 +203,7 @@ def main():
                 tokenizer,
                 new_adv_prefix_toks,
                 filter_cand=False,
-                curr_control=adv_prefix
+                curr_control=adv_prefix,
             )
 
             # Batch evaluation for all candidates with combined scoring
@@ -185,25 +214,36 @@ def main():
             count_range = max(1, max_count - min_count)
             inputs = tokenizer(candidate_texts, return_tensors="pt", padding=True, truncation=True)
             inputs = {k: v.to(device) for k, v in inputs.items()}
+
             with torch.no_grad():
                 logits = model(**inputs).logits
                 probs = torch.softmax(logits, dim=-1)
                 benign_scores = probs[:, benign_class_idx].cpu().numpy()
                 # Compute normalized loss for each candidate
-                losses = nn.CrossEntropyLoss(reduction="none")(logits, torch.zeros(logits.shape[0], device=device, dtype=torch.long))
+                losses = nn.CrossEntropyLoss(reduction="none")(
-                normalized_losses = (1.0 / (1.0 + losses.cpu().numpy()))
+                    logits,
+                    torch.zeros(logits.shape[0], device=device, dtype=torch.long),
+                )
+                normalized_losses = 1.0 / (1.0 + losses.cpu().numpy())
+
             # Compute token penalty for each candidate
-            token_penalties = [1.0 - ((tc - min_count) / count_range) if count_range > 0 else 0 for tc in token_counts]
+            token_penalties = [
+                1.0 - ((tc - min_count) / count_range) if count_range > 0 else 0
+                for tc in token_counts
+            ]
             # Compute combined score for each candidate
             combined_scores = [
-                (alpha * benign_scores[i] + (1 - alpha) * normalized_losses[i]) * (1 - token_penalty_weight + token_penalty_weight * token_penalties[i])
+                (alpha * benign_scores[i] + (1 - alpha) * normalized_losses[i])
+                * (1 - token_penalty_weight + token_penalty_weight * token_penalties[i])
                 for i in range(len(new_adv_prefix))
             ]
             idx = int(max(range(len(combined_scores)), key=lambda i: combined_scores[i]))
             adv_prefix = new_adv_prefix[idx]
 
             # Update the tokens for the next iteration
-            adv_prefix_tokens = tokenizer(adv_prefix, return_tensors="pt", add_special_tokens=False)["input_ids"][0]
+            adv_prefix_tokens = tokenizer(
+                adv_prefix, return_tensors="pt", add_special_tokens=False
+            )["input_ids"][0]
             adv_prefix_tokens = adv_prefix_tokens.to(device)
 
             # Check the current classification
@@ -216,10 +256,14 @@ def main():
                 predicted_class_id: int = logits.argmax().item()
                 benign_score: float = probs[0][benign_class_idx].item()
                 benign_percentage: float = benign_score * 100
-                malicious_score: float = probs[0][malicious_class_idx].item() if malicious_class_idx is not None else 0
+                malicious_score: float = (
+                    probs[0][malicious_class_idx].item() if malicious_class_idx is not None else 0
+                )
 
                 # Calculate combined score
-                loss: torch.Tensor = nn.CrossEntropyLoss()(logits, torch.zeros(logits.shape[0], device=device).long())
+                loss: torch.Tensor = nn.CrossEntropyLoss()(
+                    logits, torch.zeros(logits.shape[0], device=device).long()
+                )
                 normalized_loss: float = 1.0 / (1.0 + loss.item())
                 current_score: float = alpha * benign_score + (1 - alpha) * normalized_loss
 
@@ -244,9 +288,11 @@ def main():
                 if current_token_count < min_token_count:
                     min_token_count = current_token_count
 
-                print(f"Iteration {i+1}: Class={model.config.id2label[predicted_class_id]} " +
+                print(
-                        f"(benign: {benign_percentage:.2f}%, loss_norm: {normalized_loss:.4f}, " +
+                    f"Iteration {i+1}: Class={model.config.id2label[predicted_class_id]} "
-                        f"tokens: {current_token_count}, prefix: {adv_prefix})")
+                    + f"(benign: {benign_percentage:.2f}%, loss_norm: {normalized_loss:.4f}, "
+                    + f"tokens: {current_token_count}, prefix: {adv_prefix})"
+                )
 
                 if current_score > best_iteration_score:
                     # New best score, reset counter
@@ -254,46 +300,75 @@ def main():
                     iterations_without_improvement = 0
                 elif current_score >= combined_avg * improvement_threshold:
                     # Score is close enough to combined average, don't count against patience
-                    print(f"  Score within {(1-improvement_threshold)*100:.1f}% of combined average, continuing optimization")
+                    print(
+                        f"  Score within {(1-improvement_threshold)*100:.1f}% of combined average, continuing optimization"
+                    )
                     # Don't increment iterations_without_improvement
                 else:
                     # Score is significantly worse than combined average, count against patience
                     iterations_without_improvement += 1
-                    print(f"  No significant improvement for {iterations_without_improvement}/{patience} iterations")
+                    print(
+                        f"  No significant improvement for {iterations_without_improvement}/{patience} iterations"
+                    )
 
                     # If we're stagnating but not yet at early stopping threshold, try injecting educational text
-                    if iterations_without_improvement % stagnation_threshold == 0 and iterations_without_improvement < patience:
+                    if (
-                        print(f"\n  Optimization stagnating. Looking for words to improve benign rating...")
+                        iterations_without_improvement % stagnation_threshold == 0
+                        and iterations_without_improvement < patience
+                    ):
+                        print(
+                            f"\n  Optimization stagnating. Looking for words to improve benign rating..."
+                        )
 
                         # Try to find the best word to add
                         new_prefix: Optional[str]
                         improvement: float
                         new_prefix, improvement = find_best_word_to_add(
-                            model, tokenizer, injection_text, adv_prefix, text, 
+                            model,
-                            benign_class_idx, device=device, num_candidates=len(words),
+                            tokenizer,
+                            injection_text,
+                            adv_prefix,
+                            text,
+                            benign_class_idx,
+                            device=device,
+                            num_candidates=len(words),
                             token_priority=general_token_priority,  # Equal weight to token count and improvement
                         )
 
                         if new_prefix and improvement > 0:
                             # Use the optimized prefix with the best word added
                             adv_prefix = new_prefix
-                            print(f"  Applied optimized prefix with improvement of {improvement:.4f}")
+                            print(
+                                f"  Applied optimized prefix with improvement of {improvement:.4f}"
+                            )
                         else:
                             # Fall back to adding random words if no improvement found
-                            snippet: str = " ".join(get_random_words(words_to_inject, 1, token_priority=general_token_priority))
+                            snippet: str = " ".join(
+                                get_random_words(
+                                    words_to_inject,
+                                    1,
+                                    token_priority=general_token_priority,
+                                )
+                            )
 
                             # Insert the snippet at the beginning
                             adv_prefix = snippet + " " + adv_prefix
-                            print(f"  No improvement found, inserted random words at beginning: '{snippet}'")
+                            print(
+                                f"  No improvement found, inserted random words at beginning: '{snippet}'"
+                            )
 
                         # Update tokens for next iteration
-                        adv_prefix_tokens = tokenizer(adv_prefix, return_tensors="pt", add_special_tokens=False)["input_ids"][0]
+                        adv_prefix_tokens = tokenizer(
+                            adv_prefix, return_tensors="pt", add_special_tokens=False
+                        )["input_ids"][0]
                         adv_prefix_tokens = adv_prefix_tokens.to(device)
                         control_slice = slice(0, len(adv_prefix_tokens))
 
                         # Give the model time to improve with the new text by resetting best score tracking
-                        best_iteration_score = float('-inf')
+                        best_iteration_score = float("-inf")
-                        iterations_without_improvement = max(0, iterations_without_improvement - grace_period)
+                        iterations_without_improvement = max(
+                            0, iterations_without_improvement - grace_period
+                        )
                         print(f"  Reset optimization tracking to give new text time to work")
 
                 # Early stopping check - only stop if consistently no improvement
@@ -302,35 +377,53 @@ def main():
                     break
 
                 # Always track the best prefix we've seen, even if not high confidence
-                if model.config.id2label[predicted_class_id].lower() == benign_class and current_score > best_score:
+                if (
+                    model.config.id2label[predicted_class_id].lower() == benign_class
+                    and current_score > best_score
+                ):
                     best_score = current_score
                     best_prefix = adv_prefix
-                    print(f"New best benign prefix found with score: {best_score:.4f}, benign confidence: {benign_percentage:.2f}%")
+                    print(
+                        f"New best benign prefix found with score: {best_score:.4f}, benign confidence: {benign_percentage:.2f}%"
+                    )
 
                 # Check if we've found a high confidence benign classification (>95%)
-                if model.config.id2label[predicted_class_id].lower() == benign_class and benign_score > min_benign_confidence:
+                if (
+                    model.config.id2label[predicted_class_id].lower() == benign_class
+                    and benign_score > min_benign_confidence
+                ):
                     found_high_confidence_benign = True  # Set the flag
                     best_score = current_score
                     best_prefix = adv_prefix
-                    print(f"Found high confidence benign classification ({benign_percentage:.2f}%) at iteration {i+1}! Stopping optimization.")
+                    print(
+                        f"Found high confidence benign classification ({benign_percentage:.2f}%) at iteration {i+1}! Stopping optimization."
+                    )
                     break  # Stop optimizing
 
         except Exception as e:
             # print stack trace
             import traceback
+
             traceback.print_exc()
             print(f"Error in iteration {i+1}: {str(e)}")
             continue
 
     # Print final results
     if found_high_confidence_benign:
-        print(f"\n===== SUCCESS! Found a high confidence benign classification (>{min_benign_confidence*100}%) =====")
+        print(
+            f"\n===== SUCCESS! Found a high confidence benign classification (>{min_benign_confidence*100}%) ====="
+        )
 
         # Now minimize the tokens while maintaining the benign classification
         if best_prefix is not None:
             best_prefix = minimize_tokens(
-                model, tokenizer, injection_text, best_prefix, text, 
+                model,
-                benign_class_idx, device=device, 
+                tokenizer,
+                injection_text,
+                best_prefix,
+                text,
+                benign_class_idx,
+                device=device,
                 min_acceptable_benign=min_acceptable_benign,
             )
     else:
@@ -340,8 +433,13 @@ def main():
 
             # Still try to minimize tokens
             best_prefix = minimize_tokens(
-                model, tokenizer, injection_text, best_prefix, text, 
+                model,
-                benign_class_idx, device=device, 
+                tokenizer,
+                injection_text,
+                best_prefix,
+                text,
+                benign_class_idx,
+                device=device,
                 min_acceptable_benign=min_acceptable_benign,
             )
 
@@ -355,8 +453,10 @@ def main():
     probs: torch.Tensor = torch.softmax(logits, dim=-1)
     predicted_class_id: int = logits.argmax().item()
     benign_prob: float = probs[0][benign_class_idx].item() * 100
-    print(f'\nPayload is: {text}')
+    print(f"\nPayload is: {text}")
-    print(f'Payload without prefix is classified as: {model.config.id2label[predicted_class_id]} (benign probability: {benign_prob:.2f}%)')
+    print(
+        f"Payload without prefix is classified as: {model.config.id2label[predicted_class_id]} (benign probability: {benign_prob:.2f}%)"
+    )
 
     # With adversarial prefix
     full_text = injection_text + adv_prefix + text
@@ -366,15 +466,17 @@ def main():
     probs = torch.softmax(logits, dim=-1)
     predicted_class_id = logits.argmax().item()
     benign_prob = probs[0][benign_class_idx].item() * 100
-    print(f'Prefix is: {adv_prefix}')
+    print(f"Prefix is: {adv_prefix}")
-    print(f'Complete text using template is: {full_text}')
+    print(f"Complete text using template is: {full_text}")
-    print(f'Payload with prefix is classified as: {model.config.id2label[predicted_class_id]} (benign probability: {benign_prob:.2f}%)')
+    print(
+        f"Payload with prefix is classified as: {model.config.id2label[predicted_class_id]} (benign probability: {benign_prob:.2f}%)"
+    )
 
     # Try to run inference with the model's classifier
     classifier = pipeline(
         "text-classification",
         model=model_id,
-        device=0 if torch.cuda.is_available() else -1
+        device=0 if torch.cuda.is_available() else -1,
     )
 
     try:
@@ -388,8 +490,9 @@ def main():
     except Exception as e:
         print(f"Error running classifier pipeline: {str(e)}")
 
-    print(f'Adv prefix token count: {count_tokens(adv_prefix)}')
+    print(f"Adv prefix token count: {count_tokens(adv_prefix)}")
-    print(f'Total token count: {count_tokens(full_text)}')
+    print(f"Total token count: {count_tokens(full_text)}")
+
 
 if __name__ == "__main__":
     main()

utils.py

@@ -10,6 +10,7 @@ from wordsdb import WordsDatabase
 # Create a global instance of the database
 words_db = WordsDatabase()
 
+
 def find_best_word_to_add(
     model: AutoModelForSequenceClassification,
     tokenizer: AutoTokenizer,
@@ -22,7 +23,7 @@ def find_best_word_to_add(
     token_weight: float = 0.5,  # Weight for token count prioritization
     use_db: bool = True,  # Whether to use the database for word selection and tracking
     token_priority: float = 0.3,  # How much to prioritize words with fewer tokens when selecting from database
-    order_template: str = "{injection}{prefix}{text}"  # Template for ordering components
+    order_template: str = "{injection}{prefix}{text}",  # Template for ordering components
 ) -> Tuple[Optional[str], float]:
     """
     Evaluate multiple candidate words and find the one that most improves the benign score when added to the prefix.
@@ -72,12 +73,14 @@ def find_best_word_to_add(
                 limit=db_candidates_count,
                 min_uses=1,  # Only need to have been tested once
                 sort_by="combined" if token_priority > 0 else "improvement",
-                token_weight=token_priority
+                token_weight=token_priority,
             )
 
             # If we got some words from the database, use them plus some random words
             if top_words:
-                print(f"Using {len(top_words)} words from database (with token priority {token_priority}) plus {num_candidates - len(top_words)} random words")
+                print(
+                    f"Using {len(top_words)} words from database (with token priority {token_priority}) plus {num_candidates - len(top_words)} random words"
+                )
                 remaining = num_candidates - len(top_words)
                 candidates = top_words + random.choices(words, k=remaining)
             else:
@@ -103,22 +106,26 @@ def find_best_word_to_add(
                 test_prefix = adv_prefix + " " + word
             else:  # middle
                 # Find a reasonable spot to insert in the middle if possible
-                if ' ' in adv_prefix:
+                if " " in adv_prefix:
                     words_list: List[str] = adv_prefix.split()
                     middle_idx: int = len(words_list) // 2
                     words_list.insert(middle_idx, word)
-                    test_prefix = ' '.join(words_list)
+                    test_prefix = " ".join(words_list)
                 else:
                     # If no spaces, insert at midpoint of string
                     middle_idx: int = len(adv_prefix) // 2
-                    test_prefix = adv_prefix[:middle_idx] + " " + word + " " + adv_prefix[middle_idx:]
+                    test_prefix = (
+                        adv_prefix[:middle_idx] + " " + word + " " + adv_prefix[middle_idx:]
+                    )
 
-            all_candidate_prefixes.append({
+            all_candidate_prefixes.append(
-                "prefix": test_prefix,
+                {
-                "word": word,
+                    "prefix": test_prefix,
-                "position": position,
+                    "word": word,
-                "token_count": len(tokenizer.encode(test_prefix, add_special_tokens=False))
+                    "position": position,
-            })
+                    "token_count": len(tokenizer.encode(test_prefix, add_special_tokens=False)),
+                }
+            )
 
     # Prepare all candidate full texts for batch evaluation
     candidate_full_texts = [
@@ -177,7 +184,7 @@ def find_best_word_to_add(
             "tokens": token_count,
             "token_efficiency": token_efficiency,
             "combined_score": combined_score,
-            "prefix": candidate["prefix"]
+            "prefix": candidate["prefix"],
         }
 
         results.append(result)
@@ -185,11 +192,15 @@ def find_best_word_to_add(
         # Record the performance in the database if enabled
         if use_db and improvement != 0:  # Only record non-zero improvements
             words_db.record_word_performance(
-                candidate["word"], candidate["position"], benign_score, improvement, 
+                candidate["word"],
-                token_count, combined_score
+                candidate["position"],
+                benign_score,
+                improvement,
+                token_count,
+                combined_score,
             )
 
-        #print(f"Word '{candidate['word']}' at {candidate['position']}: {benign_score:.4f} (Δ: {improvement:.4f}, tokens: {token_count}, combined: {combined_score:.4f})")
+        # print(f"Word '{candidate['word']}' at {candidate['position']}: {benign_score:.4f} (Δ: {improvement:.4f}, tokens: {token_count}, combined: {combined_score:.4f})")
 
         # Only consider improvements (benign_score > baseline_score)
         if improvement > 0 and combined_score > best_combined_score:
@@ -202,7 +213,9 @@ def find_best_word_to_add(
     # Print top 5 results
     print("\nTop 5 most effective additions (based on combined score):")
     for i, result in enumerate(results[:5]):
-        print(f"{i+1}. '{result['word']}' at {result['position']}: {result['score']:.4f} (Δ: {result['improvement']:.4f}, tokens: {result['tokens']}, combined: {result['combined_score']:.4f})")
+        print(
+            f"{i+1}. '{result['word']}' at {result['position']}: {result['score']:.4f} (Δ: {result['improvement']:.4f}, tokens: {result['tokens']}, combined: {result['combined_score']:.4f})"
+        )
 
     if best_result_idx >= 0:
         best_result = all_candidate_prefixes[best_result_idx]
@@ -212,13 +225,16 @@ def find_best_word_to_add(
         best_prefix = best_result["prefix"]
 
         print(f"\nBest word to add: '{best_word}' at {best_position}")
-        print(f"Improvement: {best_improvement:.4f} (from {baseline_score:.4f} to {benign_scores[best_result_idx]:.4f})")
+        print(
+            f"Improvement: {best_improvement:.4f} (from {baseline_score:.4f} to {benign_scores[best_result_idx]:.4f})"
+        )
         print(f"New prefix: '{best_prefix}'")
         return best_prefix, best_improvement
     else:
         print("No improvement found from any candidate word")
         return None, 0
 
+
 def token_gradients_combined(
     model: AutoModelForSequenceClassification,
     input_ids: torch.Tensor,
@@ -257,12 +273,12 @@ def token_gradients_combined(
         input_ids[input_slice].shape[0],
         embed_weights.shape[0],
         device=device,
-        dtype=embed_weights.dtype
+        dtype=embed_weights.dtype,
     )
     one_hot.scatter_(
         1,
         input_ids[input_slice].unsqueeze(1),
-        torch.ones(one_hot.shape[0], 1, device=device, dtype=embed_weights.dtype)
+        torch.ones(one_hot.shape[0], 1, device=device, dtype=embed_weights.dtype),
     )
     one_hot.requires_grad_()
     input_embeds: torch.Tensor = (one_hot @ embed_weights).unsqueeze(0)
@@ -271,15 +287,18 @@ def token_gradients_combined(
     embeds: torch.Tensor = model.deberta.embeddings.word_embeddings(input_ids)
     full_embeds: torch.Tensor = torch.cat(
         [
-            embeds[:input_slice.start,:],
+            embeds[: input_slice.start, :],
             input_embeds.squeeze(),
-            embeds[input_slice.stop:,:]
+            embeds[input_slice.stop :, :],
         ],
-        dim=0)
+        dim=0,
+    )
     logits: torch.Tensor = model(inputs_embeds=full_embeds.unsqueeze(0)).logits
 
     # Combined loss: minimize malicious class (standard loss) and maximize benign class
-    standard_loss: torch.Tensor = nn.CrossEntropyLoss()(logits, torch.zeros(logits.shape[0], device=device).long())
+    standard_loss: torch.Tensor = nn.CrossEntropyLoss()(
+        logits, torch.zeros(logits.shape[0], device=device).long()
+    )
 
     # Maximize benign class probability
     log_probs: torch.Tensor = torch.log_softmax(logits, dim=1)
@@ -301,7 +320,7 @@ def analyze_token_contributions(
     benign_class_idx: int,
     device: torch.device,
     min_acceptable_benign: float = 0.6,
-    order_template: str = "{injection}{prefix}{text}"  # Template for ordering components
+    order_template: str = "{injection}{prefix}{text}",  # Template for ordering components
 ) -> str:
     """
     Simple, non-batched approach to remove as many tokens as possible while keeping
@@ -327,7 +346,9 @@ def analyze_token_contributions(
     print(f"Using threshold: {threshold:.4f}")
 
     if baseline_score < threshold:
-        print(f"Baseline score {baseline_score:.4f} already below threshold {threshold:.4f}. Stopping.")
+        print(
+            f"Baseline score {baseline_score:.4f} already below threshold {threshold:.4f}. Stopping."
+        )
         return adv_prefix
 
     current_prefix = adv_prefix
@@ -339,7 +360,7 @@ def analyze_token_contributions(
     while len(remaining_tokens) > 1:
         # Try removing each token
         best_candidate = None
-        best_score = -float('inf')
+        best_score = -float("inf")
         best_idx = -1
 
         for i in range(len(remaining_tokens)):
@@ -349,7 +370,9 @@ def analyze_token_contributions(
             candidate_prefix = tokenizer.convert_tokens_to_string(tokens_without_i)
 
             # Evaluate this candidate
-            full_text = order_template.format(injection=injection_text, prefix=candidate_prefix, text=text)
+            full_text = order_template.format(
+                injection=injection_text, prefix=candidate_prefix, text=text
+            )
 
             try:
                 inputs = tokenizer(full_text, return_tensors="pt")
@@ -376,7 +399,9 @@ def analyze_token_contributions(
             current_prefix = best_candidate
             removed_token = remaining_tokens.pop(best_idx)
             removed_tokens.append(removed_token)
-            print(f"✓ Removed token {best_idx} ('{best_token}'): new score = {best_score:.4f}, tokens left: {len(remaining_tokens)}")
+            print(
+                f"✓ Removed token {best_idx} ('{best_token}'): new score = {best_score:.4f}, tokens left: {len(remaining_tokens)}"
+            )
         else:
             print(f"Cannot remove any more tokens while staying above threshold {threshold:.4f}")
             break
@@ -403,6 +428,7 @@ def analyze_token_contributions(
 
     return current_prefix
 
+
 def minimize_tokens(
     model: AutoModelForSequenceClassification,
     tokenizer: AutoTokenizer,
@@ -426,9 +452,14 @@ def minimize_tokens(
 
     # Use only token ablation approach - systematically remove tokens that contribute least
     ablation_prefix: str = analyze_token_contributions(
-        model, tokenizer, injection_text, adv_prefix, text, 
+        model,
+        tokenizer,
+        injection_text,
+        adv_prefix,
+        text,
         benign_class_idx,
-        device=device, min_acceptable_benign=min_acceptable_benign,
+        device=device,
+        min_acceptable_benign=min_acceptable_benign,
     )
 
     # Report final token count
@@ -438,7 +469,9 @@ def minimize_tokens(
     print(f"\n===== TOKEN MINIMIZATION COMPLETE =====")
     print(f"Original token count: {original_token_count}")
     print(f"Final token count: {final_token_count}")
-    print(f"Reduction: {((original_token_count - final_token_count) / original_token_count * 100):.2f}%")
+    print(
+        f"Reduction: {((original_token_count - final_token_count) / original_token_count * 100):.2f}%"
+    )
     print(f"Final prefix: '{ablation_prefix}'")
 
     return ablation_prefix
@@ -450,10 +483,10 @@ def sample_control(
     batch_size: int,
     topk: int = 256,
     temp: float = 1,
-    not_allowed_tokens: Optional[torch.Tensor] = None
+    not_allowed_tokens: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     if not_allowed_tokens is not None:
-        grad[:, not_allowed_tokens.to(grad.device)] = float('inf')
+        grad[:, not_allowed_tokens.to(grad.device)] = float("inf")
 
     top_indices: torch.Tensor = (-grad).topk(topk, dim=1).indices
     control_toks = control_toks.to(grad.device)
@@ -467,24 +500,26 @@ def sample_control(
         0,
         len(control_toks),
         max(1, len(control_toks) / actual_batch_size),  # Ensure step is at least 1
-        device=grad.device
+        device=grad.device,
     ).type(torch.int64)
 
     # Extra safety: ensure new_token_pos is within bounds of top_indices' first dimension
     new_token_pos = torch.clamp(new_token_pos, 0, grad.shape[0] - 1)
 
     new_token_val: torch.Tensor = torch.gather(
-        top_indices[new_token_pos], 1, 
+        top_indices[new_token_pos],
-        torch.randint(0, topk, (len(new_token_pos), 1), device=grad.device)
+        1,
+        torch.randint(0, topk, (len(new_token_pos), 1), device=grad.device),
     )
 
     # Ensure we don't exceed the original batch size dimension
-    new_control_toks: torch.Tensor = original_control_toks[:len(new_token_pos)].scatter_(
+    new_control_toks: torch.Tensor = original_control_toks[: len(new_token_pos)].scatter_(
         1, new_token_pos.unsqueeze(-1), new_token_val
     )
 
     return new_control_toks
 
+
 def get_random_words(n: int = 10, min_uses: int = 0, token_priority: float = 0.3) -> List[str]:
     """
     Get a list of words to use, prioritizing words that have performed well in the past.
@@ -510,10 +545,7 @@ def get_random_words(n: int = 10, min_uses: int = 0, token_priority: float = 0.3
     else:
         # Use combined sorting with the specified token weight
         top_words = words_db.get_top_words(
-            limit=n, 
+            limit=n, min_uses=min_uses, sort_by="combined", token_weight=token_priority
-            min_uses=min_uses, 
-            sort_by="combined", 
-            token_weight=token_priority
         )
 
     # If we got enough words from the database, use them
@@ -524,6 +556,7 @@ def get_random_words(n: int = 10, min_uses: int = 0, token_priority: float = 0.3
     remaining = n - len(top_words)
     return top_words + random.choices(words, k=remaining)
 
+
 def count_tokens(text: str, model: str = "gpt-3.5") -> int:
     """Count the number of tokens in a text string using tiktoken."""
     try:
@@ -533,6 +566,7 @@ def count_tokens(text: str, model: str = "gpt-3.5") -> int:
         # Fallback to a simple approximation if tiktoken fails
         return len(text.split())
 
+
 def get_combined_score(
     model: AutoModelForSequenceClassification,
     tokenizer: AutoTokenizer,
@@ -570,7 +604,7 @@ def get_combined_score(
     int
         Index of the candidate with the best combined score.
     """
-    best_combined_score: float = float('-inf')
+    best_combined_score: float = float("-inf")
     best_idx: int = 0
 
     # Find max token count for normalization
@@ -586,8 +620,12 @@ def get_combined_score(
             probs: torch.Tensor = torch.softmax(logits, dim=-1)
 
             # Calculate loss component (lower is better)
-            loss: torch.Tensor = nn.CrossEntropyLoss()(logits, torch.zeros(logits.shape[0], device=device).long())
+            loss: torch.Tensor = nn.CrossEntropyLoss()(
-            normalized_loss: float = 1.0 / (1.0 + loss.item())  # Convert to 0-1 range where higher is better
+                logits, torch.zeros(logits.shape[0], device=device).long()
+            )
+            normalized_loss: float = 1.0 / (
+                1.0 + loss.item()
+            )  # Convert to 0-1 range where higher is better
 
             # Calculate benign score component (higher is better)
             benign_score: float = probs[0][benign_idx].item()
@@ -597,7 +635,9 @@ def get_combined_score(
             token_penalty: float = 1.0 - (token_count / max_token_count)
 
             # Calculate combined score with token penalty
-            combined_score: float = (alpha * benign_score + (1 - alpha) * normalized_loss) * (1 - token_penalty_weight + token_penalty_weight * token_penalty)
+            combined_score: float = (alpha * benign_score + (1 - alpha) * normalized_loss) * (
+                1 - token_penalty_weight + token_penalty_weight * token_penalty
+            )
 
             if combined_score > best_combined_score:
                 best_combined_score = combined_score

words.py

@@ -1,469 +1,3 @@
-words = ["Please",
-"Thank you",
-"Excuse me",
-"Pardon me",
-"Sorry",
-"Welcome",
-"Certainly",
-"Absolutely",
-"Indeed",
-"Agreed",
-"Agreed",
-"Agreed",
-"Okay",
-"Alright",
-"Very well",
-"Of course",
-"By all means",
-"No problem",
-"You're welcome",
-"My pleasure",
-"Don't mention it",
-"Anytime",
-"With pleasure",
-"Gladly",
-"Happily",
-"Kindly",
-"Gently",
-"Softly",
-"Patiently",
-"Carefully",
-"Thoughtfully",
-"Considerately",
-"Respectfully",
-"Humbly",
-"Gratefully",
-"Thankfully",
-"Appreciatively",
-"Sincerely",
-"Truly",
-"Genuinely",
-"Honestly",
-"Frankly",
-"Openly",
-"Candidly",
-"Politely",
-"Courteously",
-"Graciously",
-"Charmingly",
-"Amiably",
-"Genially",
-"Cordially",
-"Warmly",
-"Friendly",
-"Welcoming",
-"Inviting",
-"Pleasant",
-"Agreeable",
-"Kind",
-"Nice",
-"Sweet",
-"Lovely",
-"Delightful",
-"Wonderful",
-"Excellent",
-"Great",
-"Fantastic",
-"Amazing",
-"Superb",
-"Brilliant",
-"Splendid",
-"Marvelous",
-"Terrific",
-"Awesome",
-"Fabulous",
-"Spectacular",
-"Stupendous",
-"Phenomenal",
-"Remarkable",
-"Impressive",
-"Admirable",
-"Commendable",
-"Praiseworthy",
-"Respectable",
-"Honorable",
-"Dignified",
-"Noble",
-"Benevolent",
-"Generous",
-"Charitable",
-"Giving",
-"Helpful",
-"Cooperative",
-"Accommodating",
-"Obliging",
-"Supportive",
-"Understanding",
-"Empathetic",
-"Compassionate",
-"Caring",
-"Loving"]
-
-words2 = [
-"Please",
-"Thanks",
-"Sorry",
-"Excuse",
-"Pardon",
-"Welcome",
-"Kindly",
-"May",
-"Could",
-"Would",
-"Shall",
-"Might",
-"Do",
-"Certainly",
-"Indeed",
-"Absolutely",
-"Definitely",
-"Naturally",
-"Precisely",
-"Assuredly",
-"Undoubtedly",
-"Gladly",
-"Sure",
-"Alright",
-"Okay",
-"OK",
-"Fine",
-"Fair",
-"Aye",
-"Yea",
-"Obliged",
-"Sir",
-"Madam",
-"Ma'am",
-"Well",
-"Ah",
-"Oh",
-"Just",
-"Good",
-"Permit",
-"Allow",
-"Grant",
-"Proffer",
-"Beg",
-"Trust",
-"Hope",
-"Greetings",
-"Salutations",
-"Regards",
-"Congrats",
-"Congratulations",
-"Bravo",
-"Kudos",
-"Farewell",
-"Adieu",
-"Ciao",
-"Gracious",
-"Mercy",
-"Bless",
-"Pray",
-"Prithee",
-"Hark",
-"Henceforth",
-"Henceforward",
-"Hence",
-"Forsooth",
-"Respectfully",
-"Sincerely",
-"Truly",
-"Frankly",
-"Honestly",
-"Genuinely",
-"Openly",
-"Candidly",
-"Politely",
-"Courteously",
-"Graciously",
-]
-
-words3 = [
-"description",
-"manifest",
-"reddit",
-"recruit",
-"flight",
-"check",
-"position",
-"respectfully",
-"bless",
-"generator",
-"reading",
-"grave",
-"medicine",
-"paper",
-"cleaning",
-"related",
-"foul",
-"width",
-"characteristics",
-"rotate",
-"logistic",
-"named",
-"correction",
-"select",
-"consider",
-"other",
-"missing",
-"advertising",
-"named",
-"inbound",
-"rate",
-"suicide",
-"shortDescription",
-"catcher",
-"concurrent",
-"chemistry",
-"fighting",
-"complain",
-"score",
-"downloading",
-"medstation",
-"Bangkok",
-"missing",
-"weebly",
-"garnitur",
-"sporto",
-"cyclosporto",
-"LinkedIn",
-"basket",
-"nut",
-"Lifettc",
-"Collect",
-"stonk",
-"vinner",
-"rønde",
-"Collect",
-"iris",
-"Simon",
-"cleaning",
-"related",
-"သာသနာ",
-"theatre",
-"gemaak",
-"куча",
-"народ",
-"correction",
-"Bang",
-"category",
-"catcher",
-"参照",
-"separate",
-"almal",
-"Bangkok",
-"missing",
-"stock",
-"youtube",
-"attention",
-"fighting",
-"respectfully",
-"Place",
-"Upload",
-"next",
-"words",
-"Moi",
-"NAMA",
-"mandar",
-"alquiler",
-"chat",
-"Sebab",
-"Perfect",
-"distinct",
-"bots",
-"Ing",
-"falt",
-"placements",
-"sivo",
-"else",
-"はお",
-"ICA",
-"Цвет",
-"Check",
-"valid",
-"earn",
-"con",
-"villa",
-"outil",
-"Sun",
-"vertido",
-"en",
-"Dub",
-"danza",
-"Articolo",
-"Vsions",
-"Cruise",
-"Saatchara",
-"ала",
-"source",
-"ungalow",
-"TITLE",
-"gén",
-"セكية",
-"Fra",
-"英会話",
-"Verstaking",
-"Just",
-"Teacher",
-"itelji",
-"Hot",
-"Palquis",
-"enez",
-"Man",
-"Recommend",
-"YouTube",
-"attention",
-"foulo",
-"original",
-"grave",
-"May",
-"compete",
-"Metro",
-    "wacomercia",
-    "this",
-    "combat",
-    "verencolor",
-    "STAM",
-    "ilä",
-    "visit",
-    "toy",
-    "additional",
-    "在中国",
-    "cnhaben",
-    "same",
-    "including",
-    "term",
-    "注意到",
-    "position",
-    "Ingredients",
-    "classification",
-    "dimensions",
-    "REVIS",
-    "meteor",
-    "information",
-    "Term",
-    "giene",
-    "Teacher",
-    "Should",
-    "gala",
-    "부",
-    "mention",
-    "postal",
-    "foul",
-    "страница",
-    "respectfully",
-    "cutive",
-    "fighting",
-        "instrui",
-    "Songs",
-    "Christian",
-    "song",
-    "all",
-    "Мал",
-    "ozou",
-    "mus",
-    "bron",
-    "rhythm",
-    "əчитель",
-    "sis",
-    "tarra",
-    "Abdul",
-    "publish",
-    "consulta",
-    "amlustra",
-    "useful",
-    "classification",
-    "brief",
-    "Fall",
-    "amina",
-    "Carbon",
-    "bertso",
-    "Attend",
-    "licenses",
-    "sections",
-    "cidos",
-    "below",
-    "ículo",
-    "gehalt",
-    "alphabet",
-    "respective",
-    "izoh",
-    "plementar",
-    "mention",
-    "information",
-    "typeof",
-    "可见",
-    "giene",
-    "suits",
-    "sponso",
-    "gossip",
-    "Medical",
-    "er",
-    "ан",
-    "other",
-    "com",
-    "chemistry",
-    "recommend",
-    "description",
-        "vagina",
-    "注意到",
-    "album",
-    "hint",
-    "places",
-    "cart",
-    "celestial",
-    "prene",
-    "netter",
-    "yours",
-    "art",
-    "traduc",
-    "Optional",
-    "all",
-    "different",
-    "casi",
-    "Boas",
-    "estra",
-    "vis",
-    "client",
-    "ABLE",
-    "under",
-    "export",
-    "EVER",
-    "pro",
-    "exe",
-    "iliye",
-    "chemical",
-    "Дести",
-    "graphics",
-    "Bean",
-    "marques",
-    "kanta",
-    "talaga",
-    "Found",
-    "Sa",
-    "Html",
-    "transformation",
-    "extra",
-    "Wort",
-    "Male",
-    "Soda",
-    "drip",
-    "recommend",
-    "Friends",
-    "avoid",
-    "formant",
-    "type",
-    "font",
-    "UA",
-    "keyword",
-    "Term",
-    "Medical"
-]
-
-words3 = list(set(words3))
-
-
 words4 = [
     "ocular",
     "spell",
@@ -491,19 +25,33 @@ words4 = [
     "recommend",
     "better",
     "previous",
-  "Camp",
+    "Camp",
-  "Megan",
+    "Megan",
-  "Voltage",
+    "Voltage",
-  "Previous",
+    "Previous",
-  "reviews",
+    "reviews",
-  "instru",
+    "instru",
-  "Exercise",
+    "Exercise",
-  "instru",
+    "instru",
-  "Russian",
+    "Russian",
-  "vaccination",
+    "vaccination",
-  "can","dread",
+    "can",
-    'all', 'Simon', 'material', 'spray', 'Subjects', 'recess', 'position', 'contrast', 'want', 'twenty', 'dependent', 'recommend', 'read',
+    "dread",
-    'sección', 'Hospital',
+    "all",
+    "Simon",
+    "material",
+    "spray",
+    "Subjects",
+    "recess",
+    "position",
+    "contrast",
+    "want",
+    "twenty",
+    "dependent",
+    "recommend",
+    "read",
+    "sección",
+    "Hospital",
     "citation",
     "edge",
     "solid",
@@ -534,22 +82,21 @@ words4 = [
     "今年",
     "许可证号",
     "nutrition",
-  "previous",
+    "previous",
-  "additional",
+    "additional",
-  "better",
+    "better",
-  "Word",
+    "Word",
-  "leg",
+    "leg",
-  "similar",
+    "similar",
-  "anchors",
+    "anchors",
-  "civil",
+    "civil",
-  "civil",
+    "civil",
-  "are",
+    "are",
-  "bilar",
+    "bilar",
-  "calanus",
+    "calanus",
-  "ITIS",
+    "ITIS",
-  "量的",
+    "量的",
-  "1657"
+    "1657" "depression",
-    "depression",
     "Upload",
     "overseas",
     "irrigate",
@@ -622,9 +169,22 @@ words4 = [
     "question",
     "faillets",
     "耙もの",
-    "prestencil", "vine", "birds", "help", "Container", "mention",
+    "prestencil",
-    "Primary", "participation", "Maintenance", "Categories", "malaysia",
+    "vine",
-    "vascular", "editorial", "OECD", "question", "consider",
+    "birds",
+    "help",
+    "Container",
+    "mention",
+    "Primary",
+    "participation",
+    "Maintenance",
+    "Categories",
+    "malaysia",
+    "vascular",
+    "editorial",
+    "OECD",
+    "question",
+    "consider",
     "attachment",
     "information",
     "recommend",
@@ -753,8 +313,7 @@ words4 = [
     "question",
     "consider",
     "oval",
-    "preferred"
+    "preferred" "resources",
-    "resources",
     "phrases",
     "low",
     "Mark",
@@ -794,37 +353,47 @@ words4 = [
     "form",
     "similar",
     "candid",
-    'Night', 'similar', 'atelier', 'keyword', 'repository', 'maintain', 'physique', 'excessopathy', 'article', 'information', 'recommend', 'consider'
+    "Night",
-    "irish",
+    "similar",
+    "atelier",
+    "keyword",
+    "repository",
+    "maintain",
+    "physique",
+    "excessopathy",
+    "article",
+    "information",
+    "recommend",
+    "consider" "irish",
     "accessories",
     "caption",
     "pression",
-    "secteur", # French for sector
+    "secteur",  # French for sector
     "tag",
     "category",
-    "sebelum", # Indonesian for before
+    "sebelum",  # Indonesian for before
     "zoom",
-    "reibung", # German for friction
+    "reibung",  # German for friction
     "tension",
     "nutrient",
     "layer",
     "below",
     "recommend",
-    'previous',
+    "previous",
-    'Brush',
+    "Brush",
-    'write',
+    "write",
-    'some',
+    "some",
-    'needle',
+    "needle",
-    'same',
+    "same",
-    'antioxidant',
+    "antioxidant",
-    'are',
+    "are",
-    'separate',
+    "separate",
-    '注意',
+    "注意",
-    'кула',
+    "кула",
-    'лист',
+    "лист",
-    '液压',
+    "液压",
-    'gène',
+    "gène",
-    'bel'
+    "bel",
 ]
 
 words4 = list(set(words4[:230]))

wordsdb.py

@@ -2,11 +2,13 @@ import sqlite3
 from datetime import datetime
 from typing import List, Optional, Dict, Any
 
+
 class WordsDatabase:
     """
     Database to track the performance of words when added to a prefix.
     Stores word statistics and allows querying for top-performing words.
     """
+
     def __init__(self, db_path: str = "word_performance.db"):
         """Initialize the database, creating tables if they don't exist."""
         self.db_path = db_path
@@ -20,7 +22,8 @@ class WordsDatabase:
             cursor = self.conn.cursor()
 
             # Create table for word performance
-            cursor.execute('''
+            cursor.execute(
+                """
             CREATE TABLE IF NOT EXISTS word_performance (
                 id INTEGER PRIMARY KEY,
                 word TEXT NOT NULL,
@@ -31,10 +34,12 @@ class WordsDatabase:
                 combined_score REAL NOT NULL,
                 timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
             )
-            ''')
+            """
+            )
 
             # Create table for word statistics (aggregated data)
-            cursor.execute('''
+            cursor.execute(
+                """
             CREATE TABLE IF NOT EXISTS word_stats (
                 word TEXT PRIMARY KEY,
                 avg_improvement REAL NOT NULL,
@@ -45,15 +50,23 @@ class WordsDatabase:
                 best_position TEXT NOT NULL,
                 last_updated DATETIME DEFAULT CURRENT_TIMESTAMP
             )
-            ''')
+            """
+            )
 
             self.conn.commit()
             print(f"Database initialized at {self.db_path}")
         except sqlite3.Error as e:
             print(f"Database error: {e}")
 
-    def record_word_performance(self, word: str, position: str, benign_score: float, 
+    def record_word_performance(
-                               improvement: float, token_count: int, combined_score: float):
+        self,
+        word: str,
+        position: str,
+        benign_score: float,
+        improvement: float,
+        token_count: int,
+        combined_score: float,
+    ):
         """Record the performance of a word when added to a prefix."""
         if self.conn is None:
             self.initialize_db()
@@ -62,14 +75,25 @@ class WordsDatabase:
             cursor = self.conn.cursor()
 
             # Insert performance record
-            cursor.execute('''
+            cursor.execute(
+                """
             INSERT INTO word_performance 
             (word, position, benign_score, improvement, token_count, combined_score)
             VALUES (?, ?, ?, ?, ?, ?)
-            ''', (word, position, benign_score, improvement, token_count, combined_score))
+            """,
+                (
+                    word,
+                    position,
+                    benign_score,
+                    improvement,
+                    token_count,
+                    combined_score,
+                ),
+            )
 
             # Update statistics
-            cursor.execute('''
+            cursor.execute(
+                """
             INSERT INTO word_stats 
             (word, avg_improvement, max_improvement, avg_token_count, min_token_count, use_count, best_position)
             VALUES (?, ?, ?, ?, ?, 1, ?)
@@ -81,18 +105,35 @@ class WordsDatabase:
                 use_count = use_count + 1,
                 best_position = CASE WHEN ? > max_improvement THEN ? ELSE best_position END,
                 last_updated = CURRENT_TIMESTAMP
-            ''', (
+            """,
-                word, improvement, improvement, token_count, token_count, position, 
+                (
-                improvement, improvement, token_count, token_count, improvement, position
+                    word,
-            ))
+                    improvement,
+                    improvement,
+                    token_count,
+                    token_count,
+                    position,
+                    improvement,
+                    improvement,
+                    token_count,
+                    token_count,
+                    improvement,
+                    position,
+                ),
+            )
 
             self.conn.commit()
         except sqlite3.Error as e:
             print(f"Error recording word performance: {e}")
             # Still try to continue without failing
 
-    def get_top_words(self, limit: int = 20, min_uses: int = 2, sort_by: str = "improvement", 
+    def get_top_words(
-                     token_weight: float = 0.0) -> List[str]:
+        self,
+        limit: int = 20,
+        min_uses: int = 2,
+        sort_by: str = "improvement",
+        token_weight: float = 0.0,
+    ) -> List[str]:
         """
         Get the top-performing words based on selected criteria.
 
@@ -116,19 +157,25 @@ class WordsDatabase:
             # Different sorting strategies
             if sort_by == "tokens":
                 # Sort by token count (ascending) then by improvement (descending)
-                cursor.execute('''
+                cursor.execute(
+                    """
                 SELECT word FROM word_stats 
                 WHERE use_count >= ? AND avg_improvement > 0
                 ORDER BY min_token_count ASC, avg_improvement DESC
                 LIMIT ?
-                ''', (min_uses, limit))
+                """,
+                    (min_uses, limit),
+                )
             elif sort_by == "combined":
                 # Get all qualifying words with their stats
-                cursor.execute('''
+                cursor.execute(
+                    """
                 SELECT word, avg_improvement, min_token_count 
                 FROM word_stats 
                 WHERE use_count >= ? AND avg_improvement > 0
-                ''', (min_uses,))
+                """,
+                    (min_uses,),
+                )
 
                 # Calculate combined scores
                 results = cursor.fetchall()
@@ -144,8 +191,12 @@ class WordsDatabase:
                 for row in results:
                     word = row[0]
                     norm_improvement = row[1] / max_improvement if max_improvement > 0 else 0
-                    norm_tokens = 1 - (row[2] / max_tokens if max_tokens > 0 else 0)  # Invert so lower is better
+                    norm_tokens = 1 - (
-                    combined_score = (1 - token_weight) * norm_improvement + token_weight * norm_tokens
+                        row[2] / max_tokens if max_tokens > 0 else 0
+                    )  # Invert so lower is better
+                    combined_score = (
+                        1 - token_weight
+                    ) * norm_improvement + token_weight * norm_tokens
                     scored_words.append((word, combined_score))
 
                 # Sort by combined score and return top words
@@ -153,12 +204,15 @@ class WordsDatabase:
                 return [word for word, _ in scored_words[:limit]]
             else:
                 # Default: sort by improvement
-                cursor.execute('''
+                cursor.execute(
+                    """
                 SELECT word FROM word_stats 
                 WHERE use_count >= ? AND avg_improvement > 0
                 ORDER BY avg_improvement DESC
                 LIMIT ?
-                ''', (min_uses, limit))
+                """,
+                    (min_uses, limit),
+                )
 
             results = cursor.fetchall()
             return [row[0] for row in results]
@@ -173,11 +227,14 @@ class WordsDatabase:
 
         try:
             cursor = self.conn.cursor()
-            cursor.execute('''
+            cursor.execute(
+                """
             SELECT word, avg_improvement, max_improvement, avg_token_count, min_token_count, use_count, best_position
             FROM word_stats 
             WHERE word = ?
-            ''', (word,))
+            """,
+                (word,),
+            )
 
             result = cursor.fetchone()
             if result:
@@ -188,7 +245,7 @@ class WordsDatabase:
                     "avg_token_count": result[3],
                     "min_token_count": result[4],
                     "use_count": result[5],
-                    "best_position": result[6]
+                    "best_position": result[6],
                 }
             return None
         except sqlite3.Error as e: