File size: 8,546 Bytes
d7228c8 | 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 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 | """
Autoregressive Text Generator supporting Temperature, Top-K, Top-P, Penalties, and Streaming.
"""
from typing import List, Optional, Callable, Dict, Any, Union
import torch
import torch.nn as nn
import torch.nn.functional as F
from slm.model.transformer_lm import SLMForCausalLM
from slm.tokenizer.bpe import BPETokenizer
from slm.utils.logger import get_logger
logger = get_logger("slm.sampling")
class TextGenerator:
"""
Inference and sampling engine for SLMForCausalLM.
Supports Greedy, Temperature, Top-K, Top-P (Nucleus), Typical sampling,
Repetition/Frequency/Presence penalties, and streaming token callbacks.
"""
def __init__(self, model: SLMForCausalLM, tokenizer: BPETokenizer) -> None:
"""
Initializes TextGenerator.
Args:
model: Trained SLMForCausalLM model instance.
tokenizer: BPETokenizer instance for encoding prompt and decoding tokens.
"""
self.model = model
self.tokenizer = tokenizer
def _apply_penalties(
self,
logits: torch.Tensor,
generated_tokens: List[int],
repetition_penalty: float = 1.0,
frequency_penalty: float = 0.0,
presence_penalty: float = 0.0
) -> torch.Tensor:
"""
Applies repetition, frequency, and presence penalty adjustments to raw token logits.
"""
if not generated_tokens:
return logits
token_counts: Dict[int, int] = {}
for token_id in generated_tokens:
token_counts[token_id] = token_counts.get(token_id, 0) + 1
for token_id, count in token_counts.items():
if token_id >= logits.size(-1):
continue
# Repetition penalty
if repetition_penalty != 1.0:
if logits[0, token_id] < 0:
logits[0, token_id] *= repetition_penalty
else:
logits[0, token_id] /= repetition_penalty
# Frequency penalty
if frequency_penalty > 0.0:
logits[0, token_id] -= frequency_penalty * count
# Presence penalty
if presence_penalty > 0.0:
logits[0, token_id] -= presence_penalty
return logits
def _sample_next_token(
self,
logits: torch.Tensor,
temperature: float = 1.0,
top_k: int = 0,
top_p: float = 1.0,
typical_p: float = 1.0,
min_p: float = 0.0
) -> int:
"""
Samples a single token ID from output logits using specified sampling strategy.
"""
# logits shape: [1, vocab_size]
logits = logits / max(temperature, 1e-5)
# Exclude pad and unk special tokens from generation candidates
pad_id = self.tokenizer.vocab.pad_id
unk_id = self.tokenizer.vocab.unk_id
logits[..., pad_id] = float("-inf")
logits[..., unk_id] = float("-inf")
# Minimum probability filtering
if min_p > 0.0:
probs = F.softmax(logits, dim=-1)
max_prob = torch.max(probs).item()
threshold = min_p * max_prob
logits[probs < threshold] = float("-inf")
# Top-K filtering
if top_k > 0:
top_k = min(top_k, logits.size(-1))
indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
logits[indices_to_remove] = float("-inf")
# Top-P (Nucleus) filtering
if top_p < 1.0:
sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
# Remove tokens with cumulative probability above threshold
sorted_indices_to_remove = cumulative_probs > top_p
# Shift indices to keep first token above threshold
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
sorted_indices_to_remove[..., 0] = 0
indices_to_remove = sorted_indices_to_remove.scatter(
dim=-1, index=sorted_indices, src=sorted_indices_to_remove
)
logits[indices_to_remove] = float("-inf")
# Compute final probabilities
probs = F.softmax(logits, dim=-1)
# Sample or Greedy selection
if temperature <= 1e-5:
next_token_id = torch.argmax(probs, dim=-1).item()
else:
next_token_id = torch.multinomial(probs, num_samples=1).item()
return int(next_token_id)
@torch.no_grad()
def generate(
self,
prompt: str,
max_new_tokens: int = 128,
temperature: float = 0.8,
top_k: int = 40,
top_p: float = 0.9,
repetition_penalty: float = 1.1,
frequency_penalty: float = 0.0,
presence_penalty: float = 0.0,
stop_tokens: Optional[List[str]] = None,
stream_callback: Optional[Callable[[str], None]] = None,
device: str = "auto"
) -> str:
"""
Generates text continuation for prompt text string.
Args:
prompt: Text prompt string.
max_new_tokens: Maximum number of tokens to generate.
temperature: Softmax sampling temperature.
top_k: Top-K filter limit.
top_p: Nucleus sampling threshold.
repetition_penalty: Multiplicative penalty for repeated tokens.
frequency_penalty: Additive penalty per token frequency.
presence_penalty: Additive penalty for token presence.
stop_tokens: List of string stop sequences.
stream_callback: Optional callback receiving individual streamed tokens.
device: Execution device.
Returns:
Complete generated output text.
"""
self.model.eval()
dev = next(self.model.parameters()).device
# Encode prompt (prepend BOS token, do not append EOS to generation prompt)
prompt_token_ids = self.tokenizer.encode(prompt, add_special_tokens=False)
if not prompt_token_ids or prompt_token_ids[0] != self.tokenizer.vocab.bos_id:
prompt_token_ids = [self.tokenizer.vocab.bos_id] + prompt_token_ids
input_ids = torch.tensor([prompt_token_ids], dtype=torch.long, device=dev)
generated_ids: List[int] = list(prompt_token_ids)
stop_token_ids = {self.tokenizer.vocab.eos_id}
if stop_tokens:
for st in stop_tokens:
st_ids = self.tokenizer.encode(st, add_special_tokens=False)
stop_token_ids.update(st_ids)
max_ctx = self.model.config.max_seq_len
for step in range(max_new_tokens):
# Truncate context sequence if exceeding model context window
curr_input = input_ids[:, -max_ctx:]
logits, _ = self.model(curr_input)
# Focus on logits of final position: [1, vocab_size]
next_logits = logits[:, -1, :]
# Apply repetition penalties
next_logits = self._apply_penalties(
logits=next_logits,
generated_tokens=generated_ids,
repetition_penalty=repetition_penalty,
frequency_penalty=frequency_penalty,
presence_penalty=presence_penalty
)
# Sample next token
next_token = self._sample_next_token(
logits=next_logits,
temperature=temperature,
top_k=top_k,
top_p=top_p
)
generated_ids.append(next_token)
next_tensor = torch.tensor([[next_token]], dtype=torch.long, device=dev)
input_ids = torch.cat([input_ids, next_tensor], dim=1)
# Streaming callback trigger
if stream_callback is not None:
token_str = self.tokenizer.decode([next_token], skip_special_tokens=True)
stream_callback(token_str)
# Check stop condition
if next_token in stop_token_ids:
break
# Stop if generating new turn boundary ("User:" or "\nUser")
recent_text = self.tokenizer.decode(generated_ids[len(prompt_token_ids):], skip_special_tokens=True)
if "\nUser" in recent_text or "\nUser:" in recent_text or "User:" in recent_text:
break
full_output = self.tokenizer.decode(generated_ids, skip_special_tokens=True)
return full_output
|