Upload generate.py

custom_generate/generate.py

@@ -0,0 +1,101 @@
+import torch
+import random
+
+def generate(model, input_ids, generation_config=None, **kwargs):
+    print("✨ using XTC (Exclude Top Choices) generation ✨")
+
+    generation_config = generation_config or model.generation_config
+    
+    # Setup generation parameters
+    cur_length = input_ids.shape[1]
+    max_length = generation_config.max_length or cur_length + generation_config.max_new_tokens
+    
+    # Extract XTC parameters from config or kwargs, with defaults
+    # Note: You can pass these in the generate() call or set them in generation_config
+    xtc_threshold = kwargs.get("xtc_threshold", getattr(generation_config, "xtc_threshold", 0.1))
+    xtc_probability = kwargs.get("xtc_probability", getattr(generation_config, "xtc_probability", 0.0))
+    temperature = kwargs.get("temperature", getattr(generation_config, "temperature", 1.0))
+    
+    # Identify special tokens to protect (EOS is critical, Newline is preferred if known)
+    eos_token_id = generation_config.eos_token_id
+    pad_token_id = generation_config.pad_token_id
+    
+    # Try to find newline token if possible, though hard without direct tokenizer access in this scope.
+    # Users can pass `protected_token_ids` in kwargs to be specific.
+    protected_token_ids = kwargs.get("protected_token_ids", [])
+    if eos_token_id is not None:
+        protected_token_ids.append(eos_token_id)
+
+    # Basic handling for left_padding (from original example)
+    left_padding = kwargs.get("left_padding", None)
+    if left_padding is not None:
+        if not isinstance(left_padding, int) or left_padding < 0:
+            raise ValueError(f"left_padding must be an integer larger than 0, but is {left_padding}")
+        
+        pad_token = kwargs.get("pad_token", None) or pad_token_id
+        if pad_token is None:
+            raise ValueError("pad_token is not defined")
+            
+        batch_size = input_ids.shape[0]
+        pad_tensor = torch.full(size=(batch_size, left_padding), fill_value=pad_token).to(input_ids.device)
+        input_ids = torch.cat((pad_tensor, input_ids), dim=1)
+        cur_length = input_ids.shape[1]
+
+    # Sampling Loop
+    while cur_length < max_length:
+        with torch.no_grad():
+            outputs = model(input_ids)
+            next_token_logits = outputs.logits[:, -1, :]
+
+            # 1. Apply Temperature
+            if temperature != 1.0 and temperature > 0:
+                next_token_logits = next_token_logits / temperature
+
+            # 2. Apply XTC (Exclude Top Choices)
+            # Logic ported from text-generation-webui sampler_hijack.py
+            if xtc_probability > 0.0 and random.random() < xtc_probability:
+                # Calculate probabilities for sorting
+                # We sort descending to find the "Top" choices
+                sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
+                sorted_probs = sorted_logits.softmax(dim=-1)
+
+                # Identify tokens to remove
+                sorted_indices_to_remove = torch.full_like(sorted_probs, False, dtype=torch.bool)
+                
+                # XTC Logic: 
+                # If the *next* token in the sorted list is above threshold, 
+                # it means the current token is a "Top Choice" that can be excluded, 
+                # because we still have good alternatives remaining.
+                sorted_indices_to_remove[..., :-1] = sorted_probs[..., 1:] >= xtc_threshold
+
+                # Scatter back to original indices
+                indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+
+                # Safety Check: Don't remove if special tokens (EOS, etc) are targeted
+                # If any protected token is in the removal mask, we abort XTC for this step
+                should_abort = False
+                if protected_token_ids:
+                    # Check if any protected token is marked for removal
+                    # We convert list to tensor for indexing
+                    protected_tensor = torch.tensor(protected_token_ids, device=input_ids.device)
+                    # We check if any of the columns corresponding to protected tokens are True
+                    if indices_to_remove[:, protected_tensor].any():
+                        should_abort = True
+
+                if not should_abort:
+                    next_token_logits = next_token_logits.masked_fill(indices_to_remove, -float("Inf"))
+
+            # 3. Sample
+            # Using multinomial sampling (softmax + random selection)
+            probs = torch.softmax(next_token_logits, dim=-1)
+            next_tokens = torch.multinomial(probs, num_samples=1)
+
+            # Update input_ids
+            input_ids = torch.cat((input_ids, next_tokens), dim=-1)
+            cur_length += 1
+            
+            # Stop if all batch items have hit EOS (optional optimization)
+            if eos_token_id is not None and (next_tokens == eos_token_id).all():
+                break
+
+    return input_ids