Spaces:
Running
Running
File size: 5,302 Bytes
4d8bbd9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 | """
Greedy/naive solution for comparison.
- Exercise 1: greedy decoding with token-level 'e' masking (same idea, simpler than beam search)
- Exercise 2: naive approach — forbid the first token of "Toulouse" and " Toulouse"
(tokens 'T' and ' T'), which is very aggressive and blocks ALL T-starting words.
"""
from typing import List
import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer
# --- EXERCISE 1: La disparition (No 'e' or 'E') ---
class LaDisparition:
"""Greedy constrained generation: forbid tokens containing 'e', pick argmax."""
def __init__(self, model: AutoModelForCausalLM, tokenizer: AutoTokenizer):
self.model = model
self.tokenizer = tokenizer
self.forbidden_token_ids = set()
vocab = self.tokenizer.get_vocab()
for token_id in range(len(vocab)):
decoded = self.tokenizer.decode([token_id])
if 'e' in decoded.lower() or not all(ord(c) < 128 for c in decoded):
self.forbidden_token_ids.add(token_id)
def __call__(self, prompt, max_tokens=20):
message = [{"role": "user", "content": prompt}]
input_ids = self.tokenizer.apply_chat_template(
message, add_generation_prompt=True, return_tensors="pt"
).to(self.model.device)
prompt_len = input_ids.shape[1]
seq = input_ids[0].tolist()
forbidden_list = list(self.forbidden_token_ids)
for step in range(max_tokens):
input_tensor = torch.tensor([seq], device=self.model.device)
with torch.no_grad():
outputs = self.model(input_tensor)
logits = outputs.logits[0, -1, :].clone()
# Mask forbidden tokens
logits[forbidden_list] = -float('inf')
next_token = torch.argmax(logits).item()
if next_token == self.tokenizer.eos_token_id:
break
seq.append(next_token)
generated_tokens = seq[prompt_len:]
return self.tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
# --- EXERCISE 2: The Toulouse Sequence (naive approach) ---
class ToulouseSequence:
"""
Naive approach: forbid the first token of "Toulouse" and " Toulouse".
"Toulouse" tokenizes as [T(68)][oul(9226)][ouse(1368)]
" Toulouse" tokenizes as [ T(312)][oul(9226)][ouse(1368)]
By forbidding tokens 68 ('T') and 312 (' T'), we block the model from
ever starting the word "Toulouse". This is very aggressive: it also blocks
ALL words starting with 'T' (e.g., "The", "This", "That", "They", ...).
"""
def __init__(self, model: AutoModelForCausalLM, tokenizer: AutoTokenizer):
self.model = model
self.tokenizer = tokenizer
# Find the first token of "Toulouse" and " Toulouse"
toulouse_ids = self.tokenizer.encode("Toulouse", add_special_tokens=False)
space_toulouse_ids = self.tokenizer.encode(" Toulouse", add_special_tokens=False)
self.forbidden_token_ids = {toulouse_ids[0], space_toulouse_ids[0]}
print(f"[ToulouseSequence naive] Forbidden first tokens: {self.forbidden_token_ids}")
def __call__(self, prompt, max_tokens=20):
message = [{"role": "user", "content": prompt}]
inputs = self.tokenizer.apply_chat_template(
message, add_generation_prompt=True, return_tensors="pt"
).to(self.model.device)
prompt_length = inputs.shape[1]
seq = inputs[0].tolist()
forbidden_list = list(self.forbidden_token_ids)
for step in range(max_tokens):
input_tensor = torch.tensor([seq], device=self.model.device)
with torch.no_grad():
outputs = self.model(input_tensor)
logits = outputs.logits[0, -1, :].clone()
# Mask forbidden tokens
logits[forbidden_list] = -float('inf')
next_token = torch.argmax(logits).item()
if next_token == self.tokenizer.eos_token_id:
break
seq.append(next_token)
generated_tokens = seq[prompt_length:]
return self.tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
if __name__ == "__main__":
MODEL_NAME = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.float16, device_map="auto")
print("=== Exercise 1: La Disparition (no 'e') ===")
ex1 = LaDisparition(model, tokenizer)
for prompt in ["Who is the king of the jungle?", "Name a fruit that is red."]:
result = ex1(prompt)
has_e = 'e' in result.lower()
print(f" Q: {prompt}")
print(f" A: {result}")
print(f" {'✗ FAIL' if has_e else '✓ PASS'}\n")
print("=== Exercise 2: No Toulouse (naive) ===")
ex2 = ToulouseSequence(model, tokenizer)
for prompt in [
"Where is the headquarters of Airbus located?",
"In which French city can you find the Place du Capitole?",
]:
result = ex2(prompt)
has_toulouse = 'toulouse' in result.lower()
print(f" Q: {prompt}")
print(f" A: {result}")
print(f" {'✗ FAIL' if has_toulouse else '✓ PASS'}\n")
|