Beam Search

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("gpt2-medium")
tokenizer = AutoTokenizer.from_pretrained("gpt2-medium")

prompt = "Tôi yêu "
inputs = tokenizer(prompt, return_tensors="pt")

# 1. Greedy (num_beams = 1, không sample)
greedy = model.generate(
    **inputs,
    max_new_tokens=20,
    num_beams=1,
    do_sample=False,
)
print("Greedy:", tokenizer.decode(greedy[0], skip_special_tokens=True))

# 2. Beam search (num_beams = 5)
beam = model.generate(
    **inputs,
    max_new_tokens=20,
    num_beams=5,
    do_sample=False,
    length_penalty=0.8,          # 0.6–1.0 là phổ biến
    early_stopping=True,          # dừng khi đủ k beam đạt <eos>
    no_repeat_ngram_size=3,       # chặn lặp 3-gram (tránh mode collapse)
    num_return_sequences=3,       # trả về 3 beam tốt nhất
)
for i, seq in enumerate(beam):
    print(f"Beam {i}:", tokenizer.decode(seq, skip_special_tokens=True))

# 3. Nucleus sampling (top_p = 0.9)
sample = model.generate(
    **inputs,
    max_new_tokens=20,
    do_sample=True,
    top_p=0.9,
    top_k=50,
    temperature=0.8,
    num_return_sequences=3,
)
for i, seq in enumerate(sample):
    print(f"Sample {i}:", tokenizer.decode(seq, skip_special_tokens=True))

# 4. Diverse beam search
diverse = model.generate(
    **inputs,
    max_new_tokens=20,
    num_beams=6,
    num_beam_groups=3,            # chia k thành 3 nhóm
    diversity_penalty=1.0,         # hình phạt giữa các nhóm
    num_return_sequences=3,
    do_sample=False,
)
for i, seq in enumerate(diverse):
    print(f"Diverse {i}:", tokenizer.decode(seq, skip_special_tokens=True))

import torch
import torch.nn.functional as F
from typing import List, Tuple

def beam_search_decode(
    model,
    input_ids: torch.Tensor,
    beam_width: int = 5,
    max_length: int = 30,
    length_penalty: float = 0.7,
    eos_token_id: int = 2,
) -> List[Tuple[torch.Tensor, float]]:
    """
    Beam search giải mã tự viết tay.
    Trả về danh sách (sequence, score) đã sắp xếp.
    """
    device = input_ids.device

    # Mỗi beam: (sequence, cumulative_log_prob, is_done)
    beams: List[Tuple[torch.Tensor, float, bool]] = [
        (input_ids, 0.0, False)
    ]

    for step in range(max_length):
        candidates: List[Tuple[torch.Tensor, float, bool]] = []

        for seq, score, done in beams:
            # Beam đã gặp <eos> thì không mở rộng nữa.
            if done:
                candidates.append((seq, score, True))
                continue

            # Forward pass
            with torch.no_grad():
                logits = model(seq).logits[:, -1, :]  # (1, V)
                log_probs = F.log_softmax(logits, dim=-1).squeeze(0)

            # Lấy top-k token
            top_log_probs, top_indices = log_probs.topk(beam_width)

            for lp, idx in zip(top_log_probs.tolist(), top_indices.tolist()):
                new_seq = torch.cat(
                    [seq, torch.tensor([[idx]], device=device)],
                    dim=1,
                )
                new_score = score + lp
                new_done = (idx == eos_token_id)
                candidates.append((new_seq, new_score, new_done))

        # Prune: lấy top beam_width theo score sau length penalty
        def _rank(item):
            seq, sc, _ = item
            L = seq.size(1)
            return sc / (L ** length_penalty)

        candidates.sort(key=_rank, reverse=True)
        beams = candidates[:beam_width]

        # Nếu tất cả beam đã done thì dừng
        if all(done for _, _, done in beams):
            break

    # Trả về (seq, score đã chia length penalty)
    return [
        (seq, sc / (seq.size(1) ** length_penalty))
        for seq, sc, _ in beams
    ]