GPT

"""
Sinh văn bản tiếng Việt với Hugging Face Transformers.
Yêu cầu: pip install transformers torch
"""
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# ─── 1. Load model và tokenizer ───
# Dùng GPT-2 (nếu có Vietnamese version: VietAI/gpt-neo-1.3B-vietnamese-news)
model_name = "gpt2"  # hoặc "VietAI/gpt-j-6B-vietnamese-news"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
model.eval()

# ─── 2. Tokenize prompt ───
prompt = "Tôi yêu"
input_ids = tokenizer.encode(prompt, return_tensors="pt")
print(f"Token IDs: {input_ids.tolist()}")
print(f"Tokens: {tokenizer.convert_ids_to_tokens(input_ids[0])}")

# ─── 3. Generation với các strategy khác nhau ───
# 3a. Greedy decode (deterministic, lặp lại)
with torch.no_grad():
    out_greedy = model.generate(
        input_ids,
        max_length=30,
        do_sample=False,        # greedy
    )
print("Greedy:", tokenizer.decode(out_greedy[0], skip_special_tokens=True))

# 3b. Temperature + top-k + top-p (ChatGPT style)
with torch.no_grad():
    out_sampled = model.generate(
        input_ids,
        max_length=30,
        do_sample=True,
        temperature=0.7,        # cân bằng
        top_k=50,               # chỉ xét 50 token cao nhất
        top_p=0.9,              # nucleus sampling
        num_return_sequences=3, # tạo 3 biến thể
    )
for i, seq in enumerate(out_sampled):
    print(f"Sample {i+1}:", tokenizer.decode(seq, skip_special_tokens=True))

# 3c. Beam search (ổn định hơn, cho dịch/tóm tắt)
with torch.no_grad():
    out_beam = model.generate(
        input_ids,
        max_length=30,
        num_beams=5,
        early_stopping=True,
        no_repeat_ngram_size=2, # tránh lặp bigram
    )
print("Beam:", tokenizer.decode(out_beam[0], skip_special_tokens=True))

"""
Tự implement softmax sampling với temperature + top-k + top-p.
Giúp hiểu rõ chuyện gì xảy ra bên trong .generate().
"""
import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")
model.eval()

def sample_next_token(logits, temperature=1.0, top_k=50, top_p=0.9):
    """Áp dụng sampling pipeline chuẩn."""
    # Chia temperature
    logits = logits / temperature

    # Top-k: giữ k token cao nhất
    if top_k > 0:
        top_k_vals, _ = torch.topk(logits, top_k)
        thresh = top_k_vals[..., -1, None]
        logits = torch.where(logits < thresh, torch.full_like(logits, -float("inf")), logits)

    # Top-p (nucleus)
    sorted_logits, sorted_idx = torch.sort(logits, descending=True)
    sorted_probs = F.softmax(sorted_logits, dim=-1)
    cum_probs = torch.cumsum(sorted_probs, dim=-1)
    sorted_mask = cum_probs > top_p
    sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
    sorted_mask[..., 0] = False
    idx_to_remove = sorted_mask.scatter(-1, sorted_idx, sorted_mask)
    logits = logits.masked_fill(idx_to_remove, -float("inf"))

    # Softmax → sample
    probs = F.softmax(logits, dim=-1)
    next_token = torch.multinomial(probs, num_samples=1)
    return next_token

# ─── Autoregressive loop ───
prompt = "Tôi yêu"
input_ids = tokenizer.encode(prompt, return_tensors="pt")
generated = input_ids.clone()

for step in range(20):
    with torch.no_grad():
        outputs = model(generated)
        # logits của token cuối cùng
        logits = outputs.logits[:, -1, :]  # shape: [batch, vocab]

    # Sample
    next_token = sample_next_token(logits, temperature=0.8, top_k=40, top_p=0.95)
    generated = torch.cat([generated, next_token], dim=-1)

    # Dừng nếu gặp EOS
    if next_token.item() == tokenizer.eos_token_id:
        break

print(tokenizer.decode(generated[0], skip_special_tokens=True))
# Ví dụ output: "Tôi yêu em rất nhiều và mãi mãi bên em..."