| """Shakespeare Constrained Decoder — deterministic style enforcement. |
| |
| Apply this LogitsProcessor during generation to guarantee Shakespeare vocabulary. |
| This is the unjailbreakable layer — it operates on raw logits, not learned behavior. |
| """ |
| from transformers import LogitsProcessor |
| import torch |
|
|
| class ShakespeareLogitProcessor(LogitsProcessor): |
| """Boosts archaic tokens and suppresses modern tokens at every generation step.""" |
|
|
| def __init__(self, tokenizer, boost=3.0, suppress=-8.0): |
| super().__init__() |
| self.boost_ids = set() |
| self.suppress_ids = set() |
|
|
| boost_words = [ |
| "thee", "thou", "thy", "thine", "hast", "dost", "doth", "ye", |
| "hath", "art", "wilt", "shalt", "wouldst", "shouldst", "canst", |
| "didst", "prithee", "forsooth", "hark", "wherefore", "methinks", |
| "verily", "perchance", "mayhap", "alas", "alack", "anon", |
| "betwixt", "hence", "thence", "whence", "ere", "oft", "nay", |
| "aye", "yonder", "yon", "fie", "lo", "'tis", "'twas", "'twere", |
| "o'er", "e'er", "ne'er", "morn", "eve", "morrow", "quill", |
| "hearken", "beseech", "tarry", "naught", "nought", "dew", |
| "mortal", "immortal", "beauteous", "wondrous", "valiant", |
| "whilst", "unto", "thereof", "herein", "wherein", "hither", |
| "thither", "whither", |
| ] |
|
|
| suppress_words = [ |
| "AI", "chatbot", "assistant", "algorithm", "neural", |
| "GPT", "LLM", "okay", "OK", "sure", "yeah", "awesome", |
| "cool", "basically", "literally", "actually", "honestly", |
| "definitely", "absolutely", "totally", "internet", "wifi", |
| "app", "website", "download", "upload", "database", "server", |
| "API", "URL", "no problem", "happy to help", "let me know", |
| ] |
|
|
| for w in boost_words: |
| for v in [w, w.capitalize(), w.upper(), f" {w}", f" {w.capitalize()}"]: |
| self.boost_ids.update(tokenizer.encode(v, add_special_tokens=False)) |
| for w in suppress_words: |
| for v in [w, w.lower(), w.upper(), f" {w}", f" {w.lower()}"]: |
| self.suppress_ids.update(tokenizer.encode(v, add_special_tokens=False)) |
| self.suppress_ids -= self.boost_ids |
| self.boost = boost |
| self.suppress = suppress |
|
|
| def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: |
| for tid in self.boost_ids: |
| if tid < scores.shape[-1]: |
| scores[:, tid] += self.boost |
| for tid in self.suppress_ids: |
| if tid < scores.shape[-1]: |
| scores[:, tid] += self.suppress |
| return scores |
|
|
