#!/usr/bin/env python # coding=utf-8 """KenT inference: nạp prefix theo chunk song song hóa đệ quy, rồi sinh token. Quy ước: * Trục batch khi nạp prefix = các chunk liên tiếp của MỘT tài liệu. * Mỗi layer tính EMA song song theo khối, các khối/chunk được ghép bằng affine-scan; kết quả bằng đúng EMA tuần tự H_t = (1 - a) H_{t-1} + a u_t. * Sau prefix, trạng thái chỉ còn H mỗi layer -> giải mã O(1) bộ nhớ theo độ dài. Ví dụ: python infer.py --model Vi-LAB/KenT-300M-Base --prompt "Xin chào" python infer.py --model ./KenT-300M-Base --prompt-file doc.txt --chunk-size 1024 """ import argparse import math import sys from typing import Iterator, List, Optional, Tuple import torch import torch.nn.functional as F from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig # --------------------------------------------------------------------------- # # Load # --------------------------------------------------------------------------- # def load(model_id: str, device: Optional[str] = None, dtype: Optional[torch.dtype] = None): device = device or ("cuda" if torch.cuda.is_available() else "cpu") if dtype is None: # FP16 trên GPU; CPU dùng FP32 vì nhiều phép FP16 chậm/không hỗ trợ dtype = torch.float16 if device.startswith("cuda") else torch.float32 tokenizer = AutoTokenizer.from_pretrained(model_id) model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True) model = model.to(device=device, dtype=dtype).eval() gen_cfg = GenerationConfig.from_pretrained(model_id) return model, tokenizer, gen_cfg # --------------------------------------------------------------------------- # # Prefix prefill # --------------------------------------------------------------------------- # def split_into_chunks(ids: torch.Tensor, chunk_size: int, pad_id: int) -> Tuple[torch.Tensor, torch.Tensor]: """(T,) -> chunks (K, C) và mask (K, C). Padding đặt ở BÊN TRÁI chunk đầu để token cuối của chunk cuối luôn là token thật.""" t = ids.numel() if t == 0: raise ValueError("Prefix rỗng.") c = max(1, min(chunk_size, t)) k = math.ceil(t / c) pad = k * c - t chunks = F.pad(ids, (pad, 0), value=pad_id).view(k, c) mask = torch.cat([ids.new_zeros(pad), ids.new_ones(t)]).view(k, c) return chunks, mask @torch.inference_mode() def prefill( model, input_ids: torch.Tensor, chunk_size: int = 512, max_parallel_chunks: int = 16, pad_id: Optional[int] = None, cache=None, ): """Nạp prefix 1D `input_ids` (T,). Trả về (logits token cuối (V,), cache). Các chunk được xử lý song song theo nhóm tối đa `max_parallel_chunks`; giữa các nhóm, trạng thái H được truyền tiếp qua cache. Truyền `cache` có sẵn để nạp tiếp một đoạn tài liệu mới vào cùng trạng thái. """ device = next(model.parameters()).device pad_id = pad_id if pad_id is not None else (model.config.pad_token_id or 0) chunks, mask = split_into_chunks(input_ids.reshape(-1).to(device), chunk_size, pad_id) k = chunks.shape[0] out = None for s in range(0, k, max_parallel_chunks): out = model( input_ids=chunks[s : s + max_parallel_chunks], attention_mask=mask[s : s + max_parallel_chunks], cache_params=cache, use_cache=True, chunked=True, logits_to_keep=1, ) cache = out.cache_params return out.logits[0, -1], cache # --------------------------------------------------------------------------- # # Sampling # --------------------------------------------------------------------------- # def _process_logits(logits: torch.Tensor, seen: torch.Tensor, cfg: GenerationConfig) -> torch.Tensor: logits = logits.float() rp = getattr(cfg, "repetition_penalty", None) if rp and rp != 1.0: penalized = torch.where(logits < 0, logits * rp, logits / rp) logits = torch.where(seen, penalized, logits) if not cfg.do_sample: return logits temp = cfg.temperature or 1.0 if temp != 1.0: logits = logits / temp top_k = cfg.top_k or 0 if top_k > 0: kth = torch.topk(logits, min(top_k, logits.numel())).values[-1] logits = logits.masked_fill(logits < kth, float("-inf")) top_p = cfg.top_p if cfg.top_p is not None else 1.0 if top_p < 1.0: sorted_logits, idx = torch.sort(logits, descending=True) probs = torch.softmax(sorted_logits, dim=-1) remove = (torch.cumsum(probs, dim=-1) - probs) > top_p # giữ token đầu vượt ngưỡng sorted_logits = sorted_logits.masked_fill(remove, float("-inf")) logits = torch.full_like(logits, float("-inf")).scatter(0, idx, sorted_logits) return logits def _eos_set(cfg: GenerationConfig, model) -> List[int]: eos = cfg.eos_token_id if cfg.eos_token_id is not None else model.config.eos_token_id if eos is None: return [] return list(eos) if isinstance(eos, (list, tuple)) else [int(eos)] @torch.inference_mode() def stream( model, input_ids: torch.Tensor, gen_cfg: GenerationConfig, chunk_size: Optional[int] = None, max_parallel_chunks: Optional[int] = None, cache=None, ) -> Iterator[int]: """Sinh từng token: prefill chunk-song-song, rồi giải mã từng bước với trạng thái H.""" chunk_size = chunk_size or getattr(gen_cfg, "kent_prefill_chunk_size", None) or 512 max_parallel_chunks = max_parallel_chunks or getattr(gen_cfg, "kent_max_parallel_chunks", None) or 16 device = next(model.parameters()).device vocab = model.config.vocab_size ids = input_ids.reshape(-1).to(device) logits, cache = prefill(model, ids, chunk_size, max_parallel_chunks, gen_cfg.pad_token_id, cache) seen = torch.zeros(vocab, dtype=torch.bool, device=device) seen[ids] = True eos = _eos_set(gen_cfg, model) max_new = gen_cfg.max_new_tokens or 256 for _ in range(max_new): scores = _process_logits(logits, seen, gen_cfg) if gen_cfg.do_sample: nxt = torch.multinomial(torch.softmax(scores, dim=-1), 1) else: nxt = scores.argmax(dim=-1, keepdim=True) tok = int(nxt) yield tok if tok in eos: break seen[tok] = True out = model(input_ids=nxt.view(1, 1), cache_params=cache, use_cache=True, logits_to_keep=1) logits, cache = out.logits[0, -1], out.cache_params def generate(model, input_ids: torch.Tensor, gen_cfg: GenerationConfig, **kw) -> torch.Tensor: new = list(stream(model, input_ids, gen_cfg, **kw)) return torch.tensor(new, dtype=torch.long) # --------------------------------------------------------------------------- # # CLI # --------------------------------------------------------------------------- # def main(): p = argparse.ArgumentParser(description="KenT inference") p.add_argument("--model", required=True, help="repo id hoặc thư mục cục bộ") p.add_argument("--prompt", default=None) p.add_argument("--prompt-file", default=None, help="tài liệu dài, nạp bằng chunk song song") p.add_argument("--device", default=None) p.add_argument("--chunk-size", type=int, default=None) p.add_argument("--max-parallel-chunks", type=int, default=None) p.add_argument("--max-new-tokens", type=int, default=None) p.add_argument("--temperature", type=float, default=None) p.add_argument("--top-k", type=int, default=None) p.add_argument("--top-p", type=float, default=None) p.add_argument("--greedy", action="store_true") p.add_argument("--seed", type=int, default=None) args = p.parse_args() if args.seed is not None: torch.manual_seed(args.seed) if args.prompt_file: text = open(args.prompt_file, encoding="utf-8").read() elif args.prompt is not None: text = args.prompt else: p.error("Cần --prompt hoặc --prompt-file") model, tok, cfg = load(args.model, args.device) if args.greedy: cfg.do_sample = False for name in ("max_new_tokens", "temperature", "top_k", "top_p"): v = getattr(args, name) if v is not None: setattr(cfg, name, v) ids = tok(text, return_tensors="pt", add_special_tokens=True).input_ids[0] print(text, end="", flush=True) buf: List[int] = [] for t in stream(model, ids, cfg, args.chunk_size, args.max_parallel_chunks): buf.append(t) piece = tok.decode(buf, skip_special_tokens=True) if not piece.endswith("\ufffd"): # chờ đủ byte của ký tự UTF-8 sys.stdout.write(piece) sys.stdout.flush() buf = [] print() if __name__ == "__main__": main()