from typing import List, Tuple import torch import torch.nn.functional as F from transformers import AutoModelForCausalLM, AutoTokenizer # --- EXERCISE 1: La disparition (No 'e' or 'E) --- class LaDisparition: """ Generate text without ever using the letter 'e' or 'E'. For this, you must use model() directly: model(input_ids) yields logits. You need to manually adjust the logits to forbid tokens containing 'e' or 'E'. REQUIREMENT: Do NOT use model.generate(). """ def __init__(self, model: AutoModelForCausalLM, tokenizer: AutoTokenizer, debug: bool = False): self.model = model self.tokenizer = tokenizer self.debug = debug # Pre-calculate forbidden token IDs (tokens that decode to contain 'e' or 'E' or non-ASCII) # Check decoded output, not just the vocab string representation self.forbidden_token_ids = set() vocab = self.tokenizer.get_vocab() for token_id in range(len(vocab)): # Decode the token to see what it actually produces decoded = self.tokenizer.decode([token_id]) # Forbid if contains 'e'/'E' or contains non-ASCII (which might hide 'e') 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, beam_width=5): # Option 2: we use self.tokenizer.apply_chat_template to tokenize the prompt message = [{"role": "user", "content": prompt}] encoded = self.tokenizer.apply_chat_template(message, add_generation_prompt=True, return_tensors="pt") input_ids = (encoded if isinstance(encoded, torch.Tensor) else encoded["input_ids"]).to(self.model.device) prompt_len = input_ids.shape[1] # Beam search: maintain multiple hypotheses # Each hypothesis: (sequence, log_prob) beams: List[Tuple[List[int], float]] = [(input_ids[0].tolist(), 0.0)] for step in range(max_tokens): candidates = [] for seq, log_prob in beams: input_tensor = torch.tensor([seq], device=self.model.device) # Get logits from model with torch.no_grad(): outputs = self.model(input_tensor) logits = outputs.logits[0, -1, :].clone() # Create mask for forbidden tokens forbidden_mask = torch.zeros_like(logits, dtype=torch.bool) forbidden_mask[list(self.forbidden_token_ids)] = True # Convert to log probabilities log_probs = F.log_softmax(logits, dim=-1) # Ensure forbidden tokens stay at -inf in log space log_probs[forbidden_mask] = -float('inf') # Get top-k tokens for this beam, excluding -inf values top_k = min(beam_width, (~forbidden_mask).sum().item()) if top_k > 0: top_log_probs, top_indices = torch.topk(log_probs, top_k) else: # No valid tokens available, skip this beam continue for token_id, token_log_prob in zip(top_indices.tolist(), top_log_probs.tolist()): if token_id == self.tokenizer.eos_token_id: candidates.append((seq, log_prob + token_log_prob)) else: candidates.append((seq + [token_id], log_prob + token_log_prob)) # Keep top beam_width candidates by log probability candidates.sort(key=lambda x: x[1], reverse=True) beams = candidates[:beam_width] # Stop if all beams ended if all(seq[-1] == self.tokenizer.eos_token_id for seq, _ in beams): break # Debug: print all beams if self.debug: print(f"\n[DEBUG Ex1] Total beams: {len(beams)}") for i, (seq, log_prob) in enumerate(beams): decoded = self.tokenizer.decode(seq, skip_special_tokens=True) print(f" Beam {i}: log_prob={log_prob:.4f} | {decoded}") # Return the best hypothesis (only the generated part) best_seq = beams[0][0] generated_tokens = best_seq[prompt_len:] return self.tokenizer.decode(generated_tokens, skip_special_tokens=True).strip() # --- EXERCISE 2: The Toulouse Sequence --- class ToulouseSequence: """ Generate text without ever using the word 'Toulouse'. For this, you must use model() directly: model(input_ids) yields logits. We mask out all tokens that if added would lead to a prefix of "Toulouse" of length at least 4. REQUIREMENT: Do NOT use model.generate(). """ def __init__(self, model, tokenizer, debug=False): self.model = model self.tokenizer = tokenizer self.debug = debug self.forbidden_word = "Toulouse" self.min_prefix_len = 4 def _get_current_word_prefix(self, decoded_sequence: str) -> str: """Find the suffix since the last non-alphabetical character.""" last_separator_idx = -1 for i in range(len(decoded_sequence) - 1, -1, -1): if not decoded_sequence[i].isalpha(): last_separator_idx = i break if last_separator_idx != -1: return decoded_sequence[last_separator_idx + 1:] else: return decoded_sequence def _get_forbidden_mask(self, seq: List[int]) -> torch.Tensor: """ Create a mask for tokens that would create a forbidden prefix of 'Toulouse'. Returns a boolean tensor where True means the token should be forbidden. """ vocab_size = len(self.tokenizer.get_vocab()) forbidden_mask = torch.zeros(vocab_size, dtype=torch.bool, device=self.model.device) # Decode the current sequence to find the current word prefix decoded_sequence = self.tokenizer.decode(seq) current_word_prefix = self._get_current_word_prefix(decoded_sequence) # If the current word prefix is empty, we don't need to check anything yet if not current_word_prefix: return forbidden_mask # Get the token IDs for the current word prefix current_word_ids = self.tokenizer.encode(current_word_prefix, add_special_tokens=False) # Iterate over all possible next tokens for token_id in range(vocab_size): # Create a hypothetical next word by adding the candidate token hypothetical_word_ids = current_word_ids + [token_id] hypothetical_word = self.tokenizer.decode(hypothetical_word_ids) # Check if the hypothetical word is a forbidden prefix (case-insensitive) if len(hypothetical_word) >= self.min_prefix_len and \ self.forbidden_word.lower().startswith(hypothetical_word.lower()): forbidden_mask[token_id] = True return forbidden_mask def __call__(self, prompt, max_tokens=20): # Option 2: we use self.tokenizer.apply_chat_template to tokenize the prompt message = [{"role": "user", "content": prompt}] encoded = self.tokenizer.apply_chat_template(message, add_generation_prompt=True, return_tensors="pt") inputs = (encoded if isinstance(encoded, torch.Tensor) else encoded["input_ids"]).to(self.model.device) prompt_length = inputs.shape[1] # Generate tokens one by one seq = inputs[0].tolist() for step in range(max_tokens): input_tensor = torch.tensor([seq], device=self.model.device) # Get logits from model with torch.no_grad(): outputs = self.model(input_tensor) logits = outputs.logits[0, -1, :].clone() # Get forbidden mask based on current word prefix forbidden_mask = self._get_forbidden_mask(seq) # Apply the mask: set forbidden tokens to -inf logits[forbidden_mask] = float('-inf') # Greedy decoding next_token = torch.argmax(logits).item() # Stop if EOS token if next_token == self.tokenizer.eos_token_id: break seq.append(next_token) # Extract only the generated tokens (skip the input prompt tokens) generated_tokens = seq[prompt_length:] generated_text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True) return generated_text.strip() if __name__ == "__main__": # NOTE: This block is for testing only. The evaluation server provides model and tokenizer. # SETUP MODEL_NAME = "HuggingFaceTB/SmolLM2-1.7B-Instruct" tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto") la_disparition_generator = LaDisparition(model, tokenizer) print("Ex 1 (No 'e'):", la_disparition_generator("Who are you?")) toulouse_sequence_generator = ToulouseSequence(model, tokenizer, debug=True) print("Ex 2 (No 'Toulouse'):", toulouse_sequence_generator("Where is Toulouse?"))