Blame for infer.py

#!/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-1B-Base --prompt "Xin chào"
    python infer.py --model ./KenT-1B-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()