jyapayne/bypass-prompt-guard-2
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Use black to format

Browse source
This commit is contained in:
Joey Yakimowich-Payne 2025-05-12 17:47:18 -06:00
commit 06eac3eaba
4 changed files with 563 additions and 794 deletions
Download patch file Download diff file Expand all files Collapse all files

257
hacking.py
View file

@ -10,16 +10,25 @@ 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
# check if cuda is available
cuda_available: bool = torch.cuda.is_available()
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,40 +85,54 @@ 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,
default=injection_text,
help="Injection text to use in the template")
parser.add_argument("--mandatory-text", 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")
parser.add_argument(
"--injection",
type=str,
default=injection_text,
help="Injection text to use in the template",
)
parser.add_argument(
"--mandatory-text",
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()
# Update the global parameters based on command line arguments
injection_text = args.injection
text = args.mandatory_text
@ -113,33 +140,35 @@ def main():
print(f"Injection text: {injection_text}")
print(f"Mandatory text: {text}")
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
rolling_scores: List[float] = [] # List to store recent scores
top_scores: List[float] = [] # List to store top scores
# Track token counts
current_token_count: int = count_tokens(adv_prefix)
min_token_count: int = current_token_count
for i in range(max_iterations):
# 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
topk=16, # More options per token
temp=1.5, # Higher temperature for more exploration
batch_size=32, # Larger batch for more candidates
topk=16, # More options per token
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))
normalized_losses = (1.0 / (1.0 + losses.cpu().numpy()))
losses = nn.CrossEntropyLoss(reduction="none")(
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]} " +
f"(benign: {benign_percentage:.2f}%, loss_norm: {normalized_loss:.4f}, " +
f"tokens: {current_token_count}, prefix: {adv_prefix})")
print(
f"Iteration {i+1}: Class={model.config.id2label[predicted_class_id]} "
+ 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:
print(f"\n Optimization stagnating. Looking for words to improve benign rating...")
if (
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,
benign_class_idx, device=device, num_candidates=len(words),
model,
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')
iterations_without_improvement = max(0, iterations_without_improvement - grace_period)
best_iteration_score = float("-inf")
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,46 +377,69 @@ 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,
benign_class_idx, device=device,
model,
tokenizer,
injection_text,
best_prefix,
text,
benign_class_idx,
device=device,
min_acceptable_benign=min_acceptable_benign,
)
else:
print("\n===== Did not find a high confidence benign classification =====")
if best_prefix is not None:
print(f"Best prefix found had a benign score of: {best_score:.4f}")
# Still try to minimize tokens
best_prefix = minimize_tokens(
model, tokenizer, injection_text, best_prefix, text,
benign_class_idx, device=device,
model,
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'Payload without prefix is classified as: {model.config.id2label[predicted_class_id]} (benign probability: {benign_prob:.2f}%)')
print(f"\nPayload is: {text}")
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'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"Prefix is: {adv_prefix}")
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}%)"
)
# 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'Total token count: {count_tokens(full_text)}')
print(f"Adv prefix token count: {count_tokens(adv_prefix)}")
print(f"Total token count: {count_tokens(full_text)}")
if __name__ == "__main__":
main()

334
utils.py
View file

@ -10,24 +10,25 @@ from wordsdb import WordsDatabase
# Create a global instance of the database
words_db = WordsDatabase()
def find_best_word_to_add(
model: AutoModelForSequenceClassification,
tokenizer: AutoTokenizer,
injection_text: str,
adv_prefix: str,
text: str,
benign_class_idx: int,
model: AutoModelForSequenceClassification,
tokenizer: AutoTokenizer,
injection_text: str,
adv_prefix: str,
text: str,
benign_class_idx: int,
device: torch.device,
num_candidates: int = 20,
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.
Prioritizes words that result in fewer tokens while still improving the benign score.
Parameters:
-----------
model: The model to evaluate with
@ -41,14 +42,14 @@ def find_best_word_to_add(
use_db: Whether to use the database for word selection and tracking
token_priority: How much to prioritize words with fewer tokens when selecting from database
order_template: Template string for ordering components (using {injection}, {prefix}, {text})
Returns:
--------
best_word: The word that most improves the benign score
improvement: The amount of improvement in benign score
"""
print(f"\n----- TESTING {num_candidates} CANDIDATE WORDS TO ADD (BATCHED) -----")
# Get baseline benign score with current prefix
try:
full_text = order_template.format(injection=injection_text, prefix=adv_prefix, text=text)
@ -62,22 +63,24 @@ def find_best_word_to_add(
except Exception as e:
print(f"Error testing baseline: {e}")
return None, 0
# Generate candidate words to test - prioritize known good words if using database
if use_db:
# Try to get high-performing words from the database, with token count consideration
db_candidates_count = num_candidates // 2
if db_candidates_count > 0:
top_words = words_db.get_top_words(
limit=db_candidates_count,
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:
@ -88,10 +91,10 @@ def find_best_word_to_add(
else:
# Just use random words if not using the database
candidates = random.choices(words, k=num_candidates)
# Define positions to test for each word
insert_positions: List[str] = ["beginning", "middle", "end"]
# Generate all candidate prefixes - one for each word + position combination
all_candidate_prefixes = []
for word in candidates:
@ -103,33 +106,37 @@ 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:]
all_candidate_prefixes.append({
"prefix": test_prefix,
"word": word,
"position": position,
"token_count": len(tokenizer.encode(test_prefix, add_special_tokens=False))
})
test_prefix = (
adv_prefix[:middle_idx] + " " + word + " " + adv_prefix[middle_idx:]
)
all_candidate_prefixes.append(
{
"prefix": test_prefix,
"word": word,
"position": position,
"token_count": len(tokenizer.encode(test_prefix, add_special_tokens=False)),
}
)
# Prepare all candidate full texts for batch evaluation
candidate_full_texts = [
order_template.format(injection=injection_text, prefix=c["prefix"], text=text)
order_template.format(injection=injection_text, prefix=c["prefix"], text=text)
for c in all_candidate_prefixes
]
if not candidate_full_texts:
print("No candidate prefixes to evaluate")
return None, 0
# Batch inference
try:
inputs = tokenizer(candidate_full_texts, return_tensors="pt", padding=True, truncation=True)
@ -141,33 +148,33 @@ def find_best_word_to_add(
except Exception as e:
print(f"Error in batch evaluation: {e}")
return None, 0
# Calculate token counts for normalization
token_counts = [c["token_count"] for c in all_candidate_prefixes]
max_token_count = max(token_counts) if token_counts else 1
# Process the results
results = []
best_combined_score = 0
best_result_idx = -1
print(f"Running score analysis for {len(all_candidate_prefixes)} candidate prefixes")
for idx, candidate in enumerate(all_candidate_prefixes):
benign_score = benign_scores[idx]
improvement = benign_score - baseline_score
token_count = candidate["token_count"]
# Calculate token efficiency (lower token count is better)
# Normalize token count to 0-1 scale (where 1 is better = fewer tokens)
token_efficiency = 1.0 - min(1.0, token_count / max_token_count)
# Calculate combined score (weighting improvement and token efficiency)
# Only consider token efficiency if improvement is positive
combined_score = 0
if improvement > 0:
combined_score = (1 - token_weight) * improvement + token_weight * token_efficiency
# Record performance in results list
result = {
"word": candidate["word"],
@ -177,55 +184,64 @@ 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)
# 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,
token_count, combined_score
candidate["word"],
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:
best_combined_score = combined_score
best_result_idx = idx
# Sort results by combined score
results.sort(key=lambda x: x["combined_score"], reverse=True)
# 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]
best_word = best_result["word"]
best_position = best_result["position"]
best_improvement = benign_scores[best_result_idx] - baseline_score
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,
input_slice: slice,
model: AutoModelForSequenceClassification,
input_ids: torch.Tensor,
input_slice: slice,
device: torch.device,
benign_class: int = 1,
malicious_class: int = 0,
benign_class: int = 1,
malicious_class: int = 0,
alpha: float = 0.5,
) -> 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)
@ -293,75 +312,79 @@ def token_gradients_combined(
def analyze_token_contributions(
model: AutoModelForSequenceClassification,
tokenizer: AutoTokenizer,
injection_text: str,
adv_prefix: str,
text: str,
benign_class_idx: int,
model: AutoModelForSequenceClassification,
tokenizer: AutoTokenizer,
injection_text: str,
adv_prefix: str,
text: str,
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
the benign score above the minimum acceptable threshold.
"""
print("\n----- ANALYZING TOKEN CONTRIBUTIONS (NO BATCHING) -----")
# Get baseline benign score
full_text = order_template.format(injection=injection_text, prefix=adv_prefix, text=text)
inputs = tokenizer(full_text, return_tensors="pt")
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)
baseline_score = probs[0][benign_class_idx].item()
print(f"Original prefix: '{adv_prefix}'")
print(f"Original benign score: {baseline_score:.4f}")
# Use exactly the min_acceptable_benign as threshold
threshold = min_acceptable_benign
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
remaining_tokens = tokenizer.tokenize(current_prefix)
print(f"Starting with {len(remaining_tokens)} tokens")
removed_tokens = []
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)):
# Create a new candidate with this token removed
tokens_without_i = remaining_tokens.copy()
token_to_remove = tokens_without_i.pop(i)
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")
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)
score = probs[0][benign_class_idx].item()
print(f" Without token {i} ('{token_to_remove}'): score = {score:.4f}")
# If this is still above threshold and better than our current best
if score >= threshold and score > best_score:
best_candidate = candidate_prefix
@ -370,17 +393,19 @@ def analyze_token_contributions(
best_token = token_to_remove
except Exception as e:
print(f" Error evaluating without token {i}: {e}")
# If we found a valid candidate, update our prefix
if best_candidate:
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
# Final results
print("\n===== TOKEN REMOVAL COMPLETE =====")
print(f"Original prefix: '{adv_prefix}'")
@ -388,28 +413,29 @@ def analyze_token_contributions(
print(f"Removed {len(removed_tokens)} tokens: {removed_tokens}")
print(f"Original token count: {len(tokenizer.tokenize(adv_prefix))}")
print(f"Final token count: {len(remaining_tokens)}")
# Final verification
full_text = order_template.format(injection=injection_text, prefix=current_prefix, text=text)
inputs = tokenizer(full_text, return_tensors="pt")
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)
final_score = probs[0][benign_class_idx].item()
print(f"Final benign score: {final_score:.4f}")
return current_prefix
def minimize_tokens(
model: AutoModelForSequenceClassification,
tokenizer: AutoTokenizer,
injection_text: str,
adv_prefix: str,
text: str,
benign_class_idx: int,
model: AutoModelForSequenceClassification,
tokenizer: AutoTokenizer,
injection_text: str,
adv_prefix: str,
text: str,
benign_class_idx: int,
device: torch.device,
min_acceptable_benign: float = 0.6,
) -> str:
@ -417,85 +443,94 @@ def minimize_tokens(
Minimize tokens using only token contribution analysis (ablation study).
This approach systematically removes tokens that contribute least to the benign classification.
Prioritizes removing shorter tokens when they have similar impacts on benign score.
Parameters:
-----------
token_length_weight: Weight for prioritizing removal of short tokens (0-1, higher = prioritize short tokens more)
"""
print("\n===== STARTING TOKEN MINIMIZATION =====")
# 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
final_token_count: int = len(tokenizer.encode(ablation_prefix, add_special_tokens=False))
original_token_count: int = len(tokenizer.encode(adv_prefix, add_special_tokens=False))
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
def sample_control(
control_toks: torch.Tensor,
grad: torch.Tensor,
batch_size: int,
topk: int = 256,
temp: float = 1,
not_allowed_tokens: Optional[torch.Tensor] = None
control_toks: torch.Tensor,
grad: torch.Tensor,
batch_size: int,
topk: int = 256,
temp: float = 1,
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)
original_control_toks: torch.Tensor = control_toks.repeat(batch_size, 1)
# Ensure batch_size doesn't exceed the size of control_toks
actual_batch_size: int = min(batch_size, len(control_toks))
new_token_pos: torch.Tensor = torch.arange(
0,
len(control_toks),
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,
torch.randint(0, topk, (len(new_token_pos), 1), device=grad.device)
top_indices[new_token_pos],
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.
Parameters:
-----------
n: Number of words to return
min_uses: Minimum number of uses a word must have to be considered from the database
token_priority: How much to prioritize words with fewer tokens (0-1)
0 = purely improvement based, 1 = purely token count based
Returns:
--------
List of words
@ -510,20 +545,18 @@ 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,
min_uses=min_uses,
sort_by="combined",
token_weight=token_priority
limit=n, min_uses=min_uses, sort_by="combined", token_weight=token_priority
)
# If we got enough words from the database, use them
if len(top_words) >= n:
return top_words[:n]
# Otherwise, use what we got plus some random words
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,14 +566,15 @@ 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,
text: str,
candidates: List[str],
benign_idx: int,
model: AutoModelForSequenceClassification,
tokenizer: AutoTokenizer,
text: str,
candidates: List[str],
benign_idx: int,
device: torch.device,
alpha: float = 0.5,
alpha: float = 0.5,
token_penalty_weight: float = 0.1,
) -> int:
"""
@ -570,13 +604,13 @@ 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
token_counts: List[int] = [count_tokens(candidate) for candidate in candidates]
max_token_count: int = max(token_counts) if token_counts else 1
for idx, candidate in enumerate(candidates):
inputs: Dict[str, torch.Tensor] = tokenizer(candidate + text, return_tensors="pt")
# Move inputs to MPS device
@ -586,21 +620,27 @@ 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())
normalized_loss: float = 1.0 / (1.0 + loss.item()) # Convert to 0-1 range where higher is better
loss: torch.Tensor = nn.CrossEntropyLoss()(
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()
# Calculate token count penalty (normalized to 0-1, where higher is better = fewer tokens)
token_count: int = token_counts[idx]
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
best_idx = idx
return best_idx
return best_idx

609
words.py
View file

@ -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",
"Megan",
"Voltage",
"Previous",
"reviews",
"instru",
"Exercise",
"instru",
"Russian",
"vaccination",
"can","dread",
'all', 'Simon', 'material', 'spray', 'Subjects', 'recess', 'position', 'contrast', 'want', 'twenty', 'dependent', 'recommend', 'read',
'sección', 'Hospital',
"Camp",
"Megan",
"Voltage",
"Previous",
"reviews",
"instru",
"Exercise",
"instru",
"Russian",
"vaccination",
"can",
"dread",
"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",
"additional",
"better",
"Word",
"leg",
"similar",
"anchors",
"civil",
"civil",
"are",
"bilar",
"calanus",
"ITIS",
"量的",
"1657"
"depression",
"previous",
"additional",
"better",
"Word",
"leg",
"similar",
"anchors",
"civil",
"civil",
"are",
"bilar",
"calanus",
"ITIS",
"量的",
"1657" "depression",
"Upload",
"overseas",
"irrigate",
@ -622,9 +169,22 @@ words4 = [
"question",
"faillets",
"耙もの",
"prestencil", "vine", "birds", "help", "Container", "mention",
"Primary", "participation", "Maintenance", "Categories", "malaysia",
"vascular", "editorial", "OECD", "question", "consider",
"prestencil",
"vine",
"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"
"resources",
"preferred" "resources",
"phrases",
"low",
"Mark",
@ -794,37 +353,47 @@ words4 = [
"form",
"similar",
"candid",
'Night', 'similar', 'atelier', 'keyword', 'repository', 'maintain', 'physique', 'excessopathy', 'article', 'information', 'recommend', 'consider'
"irish",
"Night",
"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',
'Brush',
'write',
'some',
'needle',
'same',
'antioxidant',
'are',
'separate',
'注意',
'кула',
'лист',
'液压',
'gène',
'bel'
"previous",
"Brush",
"write",
"some",
"needle",
"same",
"antioxidant",
"are",
"separate",
"注意",
"кула",
"лист",
"液压",
"gène",
"bel",
]
words4 = list(set(words4[:230]))

157
wordsdb.py
View file

@ -2,25 +2,28 @@ 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
self.conn = None
self.initialize_db()
def initialize_db(self):
"""Create the database tables if they don't exist."""
try:
self.conn = sqlite3.connect(self.db_path)
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,31 +50,50 @@ 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,
improvement: float, token_count: int, combined_score: float):
def record_word_performance(
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()
try:
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,104 +105,137 @@ 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",
token_weight: float = 0.0) -> List[str]:
def get_top_words(
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.
Parameters:
-----------
limit: Maximum number of words to return
min_uses: Minimum number of uses a word must have to be considered
sort_by: How to sort the results - options: "improvement", "tokens", "combined"
token_weight: When sort_by="combined", weight for token count vs improvement (0-1)
Returns:
--------
List of words matching the criteria
"""
if self.conn is None:
self.initialize_db()
try:
cursor = self.conn.cursor()
# 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()
if not results:
return []
# Normalize values
max_improvement = max(row[1] for row in results)
max_tokens = max(row[2] for row in results)
# Calculate combined score for each word
scored_words = []
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
combined_score = (1 - token_weight) * norm_improvement + token_weight * norm_tokens
norm_tokens = 1 - (
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
scored_words.sort(key=lambda x: x[1], reverse=True)
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]
except sqlite3.Error as e:
print(f"Error getting top words: {e}")
return []
def get_word_stats(self, word: str) -> Optional[Dict[str, Any]]:
"""Get statistics for a specific word."""
if self.conn is None:
self.initialize_db()
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:
return {
@ -188,15 +245,15 @@ 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:
print(f"Error getting word stats: {e}")
return None
def close(self):
"""Close the database connection."""
if self.conn:
self.conn.close()
self.conn = None
self.conn = None