| from collections.abc import Callable, Sequence |
| from dataclasses import dataclass, field |
| import torch |
|
|
| VERSION = "dev3.0.0" |
|
|
| TokenConstraint = Callable[[list[int]], Sequence[int] | None] |
|
|
| @dataclass(slots=True) |
| class GenerationConfig: |
| strategy: str = "auto" |
| max_new_tokens: int = 500 |
| beam_size: int = 5 |
| stable_steps: int | None = 8 |
| temperature: float = 0.8 |
| top_k: int = 0 |
| top_p: float = 0.95 |
| repetition_penalty: float = 1.2 |
| no_repeat_ngram_size: int = 3 |
| length_penalty: float = 1.0 |
| lowercase: bool = True |
| seed: int | None = None |
| stop_sequences: list[list[int]] = field(default_factory=list) |
|
|
| @dataclass(slots=True) |
| class GenerationResult: |
| text: str |
| token_ids: list[int] |
| finish_reason: str |
| prompt_tokens: int |
| completion_tokens: int |
| decode_steps: int |
| completed_beams: int |
|
|
| def _special_ids(tokenizer): |
| ids = { |
| "bos": tokenizer.piece_to_id("[BOS]"), |
| "eos": tokenizer.piece_to_id("[EOS]"), |
| "im_start": tokenizer.piece_to_id("<|im_start|>"), |
| "im_end": tokenizer.piece_to_id("<|im_end|>"), |
| } |
| if any(token_id < 0 for token_id in ids.values()): |
| raise ValueError(f"Tokenizer is missing required special tokens: {ids}") |
| return ids |
|
|
| def _normalise_messages(user_input, system=None, instruction=None, history=None, messages=None): |
| if system is not None and instruction is not None: |
| raise ValueError("Use either system or instruction, not both") |
| system = system if system is not None else instruction |
| if messages is not None: |
| if user_input not in (None, "") or system is not None or history: |
| raise ValueError("Use either messages or user_input/instruction/history, not both") |
| source = list(messages) |
| else: |
| source = [] |
| if system: |
| source.append({"role": "system", "content": system}) |
| for turn in history or []: |
| if "role" in turn: |
| source.append(turn) |
| continue |
| if turn.get("user"): |
| source.append({"role": "user", "content": turn["user"]}) |
| if turn.get("model"): |
| source.append({"role": "model", "content": turn["model"]}) |
| if user_input is not None: |
| source.append({"role": "user", "content": user_input}) |
|
|
| result = [] |
| role_map = {"assistant": "model", "model": "model", "user": "user", "system": "system"} |
| for message in source: |
| role = role_map.get(str(message.get("role", "")).lower()) |
| content = str(message.get("content", "")).strip() |
| if role is None or not content: |
| raise ValueError(f"Invalid chat message: {message!r}") |
| result.append({"role": role, "content": content}) |
| if not result or result[-1]["role"] != "user": |
| raise ValueError("The final input message must have role='user'") |
| return result |
|
|
| def _build_input(user_input, tokenizer, *, max_seq_len, config, system=None, instruction=None, history=None, messages=None): |
| special = _special_ids(tokenizer) |
| normalised = _normalise_messages(user_input, system, instruction, history, messages) |
|
|
| def encode_message(message): |
| text = f"{message['role']}\n{message['content']}" |
| if config.lowercase: |
| text = text.lower() |
| return [special["im_start"]] + tokenizer.encode(text, out_type=int) + [special["im_end"]] |
|
|
| segments = [(message["role"], encode_message(message)) for message in normalised] |
| assistant_prefix = [special["im_start"]] + tokenizer.encode("model\n", out_type=int) |
| prompt_budget = max_seq_len - config.max_new_tokens |
| required = 1 + len(assistant_prefix) |
| if prompt_budget < required: |
| raise ValueError("max_new_tokens leaves no room for a prompt") |
|
|
| system_segments = [segment for role, segment in segments if role == "system"] |
| latest = segments[-1][1] |
| fixed_len = required + len(latest) + sum(map(len, system_segments)) |
| if fixed_len > prompt_budget: |
| raise ValueError( |
| f"System + latest user message need {fixed_len} tokens, " |
| f"but prompt budget is {prompt_budget}" |
| ) |
|
|
| middle = [(role, segment) for role, segment in segments[:-1] if role != "system"] |
| turn_groups = [] |
| index = 0 |
| while index < len(middle): |
| role, segment = middle[index] |
| group = [segment] |
| if role == "user" and index + 1 < len(middle) and middle[index + 1][0] == "model": |
| group.append(middle[index + 1][1]) |
| index += 1 |
| turn_groups.append(group) |
| index += 1 |
|
|
| kept_groups = [] |
| used = fixed_len |
| for group in reversed(turn_groups): |
| group_len = sum(map(len, group)) |
| if used + group_len <= prompt_budget: |
| kept_groups.append(group) |
| used += group_len |
| kept_groups.reverse() |
|
|
| input_ids = [special["bos"]] |
| for segment in system_segments: |
| input_ids.extend(segment) |
| for group in kept_groups: |
| for segment in group: |
| input_ids.extend(segment) |
| input_ids.extend(latest) |
| input_ids.extend(assistant_prefix) |
| return input_ids, special |
|
|
| def validate_generation_config(config): |
| """Báo ValueError nếu giá trị của ``config`` nằm ngoài miền hợp lệ; gọi được trước khi nạp model.""" |
| if config.strategy not in {"auto", "beam", "greedy", "sample"}: |
| raise ValueError("strategy must be auto, beam, greedy, or sample") |
| if config.max_new_tokens < 1: |
| raise ValueError("max_new_tokens must be >= 1") |
| if config.beam_size < 1: |
| raise ValueError("beam_size must be >= 1") |
| if config.stable_steps is not None and config.stable_steps < 1: |
| raise ValueError("stable_steps must be >= 1 or None") |
| if config.temperature <= 0: |
| raise ValueError("temperature must be > 0") |
| if config.top_k < 0: |
| raise ValueError("top_k must be >= 0") |
| if not 0 < config.top_p <= 1: |
| raise ValueError("top_p must be in (0, 1]") |
| if config.repetition_penalty <= 0: |
| raise ValueError("repetition_penalty must be > 0") |
| if config.no_repeat_ngram_size < 0: |
| raise ValueError("no_repeat_ngram_size must be >= 0") |
| if config.length_penalty < 0: |
| raise ValueError("length_penalty must be >= 0") |
|
|
| def _apply_repetition_penalty(logits, generated, penalty): |
| if penalty == 1.0 or not generated: |
| return logits |
| logits = logits.clone() |
| token_ids = torch.tensor(list(set(generated)), dtype=torch.long, device=logits.device) |
| selected = logits.index_select(-1, token_ids) |
| selected = torch.where(selected < 0, selected * penalty, selected / penalty) |
| logits.scatter_(-1, token_ids, selected) |
| return logits |
|
|
| def _banned_ngram_tokens(generated, n): |
| if n <= 0 or len(generated) < n - 1: |
| return set() |
| if n == 1: |
| return set(generated) |
| prefix = tuple(generated[-(n - 1):]) |
| return { |
| generated[index + n - 1] |
| for index in range(len(generated) - n + 1) |
| if tuple(generated[index:index + n - 1]) == prefix |
| } |
|
|
| def _is_stopped(generated, stop_ids, stop_sequences): |
| if generated and generated[-1] in stop_ids: |
| return True |
| return any( |
| sequence and len(generated) >= len(sequence) and generated[-len(sequence):] == sequence |
| for sequence in stop_sequences |
| ) |
|
|
| def _decode(generated, tokenizer, stop_ids, stop_sequences): |
| output = list(generated) |
| for sequence in stop_sequences: |
| if sequence and len(output) >= len(sequence) and output[-len(sequence):] == sequence: |
| del output[-len(sequence):] |
| break |
| while output and output[-1] in stop_ids: |
| output.pop() |
| return tokenizer.decode(output) |
|
|
| def _beam_score(candidate, length_penalty): |
| length = max(len(candidate["generated"]), 1) |
| return candidate["log_prob"] / (length ** length_penalty) |
|
|
| def _process_logits(logits, generated, config, constraint=None): |
| logits = _apply_repetition_penalty(logits, generated, config.repetition_penalty) |
| banned = _banned_ngram_tokens(generated, config.no_repeat_ngram_size) |
| if banned: |
| logits = logits.clone() |
| logits[list(banned)] = -torch.inf |
| if constraint is not None: |
| allowed = constraint(generated) |
| if allowed is not None: |
| allowed = list(allowed) |
| if not allowed: |
| raise ValueError("Token constraint returned an empty allowed set") |
| constrained = torch.full_like(logits, -torch.inf) |
| constrained[allowed] = logits[allowed] |
| logits = constrained |
| return logits |
|
|
| def _filter_sampling_logits(logits, config): |
| logits = logits / config.temperature |
| if config.top_k > 0: |
| threshold = torch.topk(logits, min(config.top_k, logits.numel())).values[-1] |
| logits = logits.masked_fill(logits < threshold, -torch.inf) |
| if config.top_p < 1: |
| sorted_logits, sorted_indices = torch.sort(logits, descending=True) |
| probabilities = torch.softmax(sorted_logits, dim=-1) |
| remove = probabilities.cumsum(dim=-1) > config.top_p |
| remove[1:] = remove[:-1].clone() |
| remove[0] = False |
| logits = logits.clone() |
| logits[sorted_indices[remove]] = -torch.inf |
| return logits |
|
|
| def _single_path_generate(model, input_ids, tokenizer, special, config, strategy, constraint): |
| device = model.lm_head.weight.device |
| stop_ids = {special["eos"], special["im_end"]} |
| stop_sequences = [list(sequence) for sequence in config.stop_sequences] |
| prompt_len = len(input_ids) |
| cache = model.init_cache(1, min(model.max_seq_len, prompt_len + config.max_new_tokens), device) |
| prompt = torch.tensor([input_ids], dtype=torch.long, device=device) |
| generator = None |
| if config.seed is not None: |
| generator = torch.Generator(device=device).manual_seed(config.seed) |
|
|
| with torch.inference_mode(): |
| logits, _ = model.prefill(prompt, kv_cache=cache) |
| generated = [] |
| finish_reason = "length" |
| decode_steps = 0 |
| for step in range(config.max_new_tokens): |
| row = _process_logits(logits[0], generated, config, constraint) |
| if strategy == "greedy": |
| token = int(torch.argmax(row).item()) |
| else: |
| row = _filter_sampling_logits(row, config) |
| token = int(torch.multinomial(torch.softmax(row, dim=-1), 1, generator=generator).item()) |
| generated.append(token) |
| if _is_stopped(generated, stop_ids, stop_sequences): |
| finish_reason = "stop" |
| break |
| if step + 1 < config.max_new_tokens: |
| last = torch.tensor([[token]], dtype=torch.long, device=device) |
| logits = model.decode_step(last, cache, prompt_len + step) |
| decode_steps += 1 |
|
|
| return GenerationResult( |
| text=_decode(generated, tokenizer, stop_ids, stop_sequences), |
| token_ids=generated, |
| finish_reason=finish_reason, |
| prompt_tokens=prompt_len, |
| completion_tokens=len(generated), |
| decode_steps=decode_steps, |
| completed_beams=0, |
| ) |
|
|
| def _beam_generate(model, input_ids, tokenizer, special, config): |
| device = model.lm_head.weight.device |
| stop_ids = {special["eos"], special["im_end"]} |
| stop_sequences = [list(sequence) for sequence in config.stop_sequences] |
| prompt_len = len(input_ids) |
| cache = model.init_cache( |
| config.beam_size, |
| min(model.max_seq_len, prompt_len + config.max_new_tokens), |
| device, |
| ) |
| prompt = torch.tensor([input_ids], dtype=torch.long, device=device) |
|
|
| with torch.inference_mode(): |
| first_logits, present = model.prefill(prompt, kv_cache=None) |
| for layer, (key, value) in enumerate(present): |
| cache[layer][0][:, :, :prompt_len].copy_(key.expand(config.beam_size, -1, -1, -1)) |
| cache[layer][1][:, :, :prompt_len].copy_(value.expand(config.beam_size, -1, -1, -1)) |
|
|
| first_log_probs = torch.log_softmax(first_logits[0], dim=-1) |
| top_log_probs, top_tokens = torch.topk( |
| first_log_probs, |
| min(config.beam_size, first_log_probs.numel()), |
| ) |
| alive, completed = [], [] |
| for log_prob, token in zip(top_log_probs.tolist(), top_tokens.tolist()): |
| candidate = {"generated": [int(token)], "log_prob": float(log_prob), "source": 0} |
| target = completed if _is_stopped(candidate["generated"], stop_ids, stop_sequences) else alive |
| target.append(candidate) |
|
|
| cache_len = prompt_len |
| decode_steps = 0 |
| stable_count = 0 |
| last_best_completed = None |
| finish_reason = "length" |
|
|
| while alive and max(len(item["generated"]) for item in alive) < config.max_new_tokens: |
| last_tokens = torch.tensor( |
| [[candidate["generated"][-1]] for candidate in alive], |
| dtype=torch.long, |
| device=device, |
| ) |
| logits = model.decode_step(last_tokens, cache, cache_len) |
| cache_len += 1 |
| decode_steps += 1 |
|
|
| candidates = [] |
| for beam_index, candidate in enumerate(alive): |
| row = _process_logits(logits[beam_index], candidate["generated"], config) |
| log_probs = torch.log_softmax(row, dim=-1) |
| values, tokens = torch.topk(log_probs, min(config.beam_size, log_probs.numel())) |
| for value, token in zip(values.tolist(), tokens.tolist()): |
| candidates.append({ |
| "generated": candidate["generated"] + [int(token)], |
| "log_prob": candidate["log_prob"] + float(value), |
| "source": beam_index, |
| }) |
|
|
| candidates.sort( |
| key=lambda item: _beam_score(item, config.length_penalty), |
| reverse=True, |
| ) |
| next_alive = [] |
| for candidate in candidates: |
| if _is_stopped(candidate["generated"], stop_ids, stop_sequences): |
| completed.append(candidate) |
| elif len(next_alive) < config.beam_size: |
| next_alive.append(candidate) |
| if len(next_alive) >= config.beam_size and len(completed) >= config.beam_size: |
| break |
|
|
| if completed: |
| best_completed_score = max( |
| _beam_score(item, config.length_penalty) for item in completed |
| ) |
| if last_best_completed is None or best_completed_score > last_best_completed + 1e-12: |
| last_best_completed = best_completed_score |
| stable_count = 0 |
| else: |
| stable_count += 1 |
|
|
| if ( |
| config.stable_steps is not None |
| and len(completed) >= config.beam_size |
| and stable_count >= config.stable_steps |
| ): |
| best_alive_score = max( |
| (_beam_score(item, config.length_penalty) for item in next_alive), |
| default=float("-inf"), |
| ) |
| if best_completed_score >= best_alive_score: |
| finish_reason = "stable_beam" |
| alive = next_alive |
| break |
|
|
| if not next_alive: |
| finish_reason = "all_beams_finished" |
| break |
|
|
| sources_list = [candidate["source"] for candidate in next_alive] |
| if sources_list != list(range(len(sources_list))): |
| sources = torch.tensor(sources_list, dtype=torch.long, device=device) |
| for key, value in cache: |
| selected_key = key.index_select(0, sources)[:, :, :cache_len].clone() |
| selected_value = value.index_select(0, sources)[:, :, :cache_len].clone() |
| key[:len(next_alive), :, :cache_len].copy_(selected_key) |
| value[:len(next_alive), :, :cache_len].copy_(selected_value) |
| alive = next_alive |
|
|
| pool = completed + alive |
| if not pool: |
| return GenerationResult("", [], "empty", prompt_len, 0, decode_steps, len(completed)) |
| best = max(pool, key=lambda item: _beam_score(item, config.length_penalty)) |
| return GenerationResult( |
| text=_decode(best["generated"], tokenizer, stop_ids, stop_sequences), |
| token_ids=list(best["generated"]), |
| finish_reason=finish_reason, |
| prompt_tokens=prompt_len, |
| completion_tokens=len(best["generated"]), |
| decode_steps=decode_steps, |
| completed_beams=len(completed), |
| ) |
|
|
| def generate(model, user_input=None, tokenizer=None, *, messages=None, system=None, instruction=None, history=None, config=None, constraint: TokenConstraint | None = None, return_result=False): |
| """Generate a response with beam search, greedy decoding, or sampling.""" |
| if tokenizer is None: |
| raise ValueError("tokenizer is required") |
| config = config or GenerationConfig() |
| validate_generation_config(config) |
| strategy = config.strategy |
| if strategy == "auto": |
| strategy = "greedy" if constraint is not None else "beam" |
| if strategy == "beam" and constraint is not None: |
| raise ValueError("Beam strategy does not support token constraints") |
|
|
| input_ids, special = _build_input( |
| user_input, |
| tokenizer, |
| max_seq_len=model.max_seq_len, |
| config=config, |
| system=system, |
| instruction=instruction, |
| history=history, |
| messages=messages, |
| ) |
| if strategy == "beam": |
| result = _beam_generate(model, input_ids, tokenizer, special, config) |
| else: |
| result = _single_path_generate( |
| model, |
| input_ids, |
| tokenizer, |
| special, |
| config, |
| strategy, |
| constraint, |
| ) |
| return result if return_result else result.text |
|
|
