MiniGPT — 从代码理解 Transformer

import json
import argparse
from collections import Counter

def build_vocab(data_path: str, output_path: str):
    counter = Counter()

    with open(data_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue

            try:
                item = json.loads(line)
                # 收集 question 和 answer 中的所有字符
                counter.update(item["question"])
                counter.update(item["answer"])
            except Exception as e:
                print(f"Warning: skip invalid line: {line[:50]}... | Error: {e}")

    # 获取所有唯一字符
    chars = sorted(counter.keys())

    # 添加特殊 token
    special_tokens = ["<pad>", "<unk>", "<sep>"]

    # 构建 word2id：先放特殊 token，再放字符（顺序固定便于复现）
    word2id = {token: i for i, token in enumerate(special_tokens)}
    for char in chars:
        if char not in special_tokens:  # 防御：跳过特殊 token
            word2id[char] = len(word2id)  # 自动递增

    # 构建 id2word
    id2word = {i: token for token, i in word2id.items()}

    vocab = {"word2id": word2id, "id2word": id2word}

    # 保存 vocab
    with open(output_path, "w", encoding="utf-8") as f:
        json.dump(vocab, f, ensure_ascii=False, indent=4)

    print(f"✅ Vocabulary built and saved to {output_path}")
    print(f"   Total tokens: {len(word2id)}")
    print(f"   Special tokens: {special_tokens}")
    print(f"   Sample chars: {chars[:10]}...")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Build vocabulary from QA dataset.")
    parser.add_argument(
        "--data", type=str, required=True, help="Path to training data (JSONL format)"
    )
    parser.add_argument(
        "--output", type=str, default="vocab.json", help="Output vocab file path"
    )
    args = parser.parse_args()

    build_vocab(args.data, args.output)

import json

class Tokenizer:
    def __init__(self, vocab_path: str):
        with open(vocab_path, "r", encoding="utf-8") as f:
            vocab = json.load(f)

        self.word2id = vocab["word2id"]
        self.id2word = {int(k): v for k, v in vocab["id2word"].items()}

        # 固定特殊 token ID
        self.pad_token_id = self.word2id["<pad>"]
        self.unk_token_id = self.word2id["<unk>"]
        self.sep_token_id = self.word2id["<sep>"]

    def encode(
        self,
        question: str,
        answer: str,
        max_length: int = 128,
        pad_to_max_length: bool = True,
    ):
        """将问答对编码为 token ID 序列。"""
        tokens = []

        # encode question
        for char in question:
            tokens.append(self.word2id.get(char, self.unk_token_id))
        tokens.append(self.sep_token_id)  # 添加分隔符

        # encode answer
        if answer is not None:
            for char in answer:
                tokens.append(self.word2id.get(char, self.unk_token_id))

            tokens.append(self.sep_token_id)

        # 构建 attention mask（1=真实 token，0=padding）
        attn_mask = [1] * len(tokens)

        # 截断或填充
        if pad_to_max_length:
            if len(tokens) > max_length:
                # 截断（保留开头）
                tokens = tokens[:max_length]
                attn_mask = attn_mask[:max_length]
            else:
                # 填充
                pad_len = max_length - len(tokens)
                tokens.extend([self.pad_token_id] * pad_len)
                attn_mask.extend([0] * pad_len)

        return tokens, attn_mask

    def decode(self, ids):
        """将 token ID 列表解码为原始文本（跳过 <pad>）。"""
        return "".join(
            self.id2word[i] for i in ids if i != self.pad_token_id  # 跳过填充符
        )

    def get_vocab_size(self):
        return len(self.id2word)

if __name__ == "__main__":
    question = "你好，最近怎么样？"
    answer = "我很好，谢谢！"

    tokenizer = Tokenizer("./data/vocab.json")

    input_ids, attn_mask = tokenizer.encode(question, answer, max_length=32)
    print(input_ids)
    print(tokenizer.decode(input_ids))

# split_data.py
import json
import argparse
import random
from pathlib import Path

def split_jsonl_data(input_path: str, train_ratio: float = 0.9, seed: int = 42):
    """
    将 JSONL 格式的 QA 数据集划分为 train.jsonl 和 val.jsonl。
    """
    input_path = Path(input_path)
    if not input_path.exists():
        raise FileNotFoundError(f"Input file not found: {input_path}")

    # 读取所有有效行
    data = []
    with open(input_path, "r", encoding="utf-8") as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                item = json.loads(line)
                if "question" in item and "answer" in item:
                    data.append(line)  # 保留原始字符串，避免格式变化
                else:
                    print(
                        f"⚠️  Warning: Line {line_num} missing 'question' or 'answer', skipped."
                    )
            except json.JSONDecodeError:
                print(f"⚠️  Warning: Line {line_num} is invalid JSON, skipped.")

    if not data:
        raise ValueError("No valid data found!")

    # 打乱并划分
    random.seed(seed)
    random.shuffle(data)
    n_train = int(len(data) * train_ratio)

    train_data = data[:n_train]
    val_data = data[n_train:]

    # 输出路径
    output_dir = input_path.parent
    train_path = output_dir / "train.jsonl"
    val_path = output_dir / "val.jsonl"

    # 写入文件
    with open(train_path, "w", encoding="utf-8") as f:
        f.write("\n".join(train_data) + "\n")
    with open(val_path, "w", encoding="utf-8") as f:
        f.write("\n".join(val_data) + "\n")

    print("✅ Split completed!")
    print(f"   Total samples: {len(data)}")
    print(f"   Train: {len(train_data)} → {train_path}")
    print(f"   Val:   {len(val_data)} → {val_path}")
    print(f"   Train ratio: {train_ratio:.1%}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Split QA dataset into train/val sets."
    )
    parser.add_argument(
        "--input",
        type=str,
        required=True,
        help="Path to original JSONL dataset (e.g., data/all.jsonl)",
    )
    parser.add_argument(
        "--train_ratio",
        type=float,
        default=0.9,
        help="Proportion of data to use for training (default: 0.9)",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="Random seed for reproducibility (default: 42)",
    )
    args = parser.parse_args()

    split_jsonl_data(args.input, args.train_ratio, args.seed)

import json
import argparse

import numpy as np

from minigpt.tokenizer import Tokenizer


def main():
    parser = argparse.ArgumentParser(description="Token length statistics for a JSONL QA dataset.")
    parser.add_argument("--train_data", type=str, required=True, help="Path to train.jsonl")
    parser.add_argument("--vocab_path", type=str, required=True, help="Path to vocab.json")
    parser.add_argument("--max_length", type=int, default=140)
    args = parser.parse_args()

    tokenizer = Tokenizer(args.vocab_path)

    lengths, q_lens, a_lens = [], [], []
    with open(args.train_data, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            item = json.loads(line)
            q, a = item["question"], item["answer"]
            q_lens.append(len(q))
            a_lens.append(len(a))
            # 完整序列 = question + <sep> + answer + <sep>
            lengths.append(len(q) + 1 + len(a) + 1)

    lengths = np.array(lengths)
    q_lens = np.array(q_lens)
    a_lens = np.array(a_lens)

    print(f"samples: {len(lengths)}")
    print(f"full-seq length  mean={lengths.mean():.1f}  p50={np.percentile(lengths,50):.0f}  "
          f"p90={np.percentile(lengths,90):.0f}  p95={np.percentile(lengths,95):.0f}  "
          f"p99={np.percentile(lengths,99):.0f}  max={lengths.max()}")
    print(f"question length  mean={q_lens.mean():.1f}  p95={np.percentile(q_lens,95):.0f}  max={q_lens.max()}")
    print(f"answer   length  mean={a_lens.mean():.1f}  p95={np.percentile(a_lens,95):.0f}  max={a_lens.max()}")
    for thr in (128, 140, 256):
        print(f"  <= {thr:3d}: {np.mean(lengths <= thr) * 100:.2f}%")
    print(f"truncated at max_length={args.max_length}: {np.mean(lengths > args.max_length) * 100:.2f}%")


if __name__ == "__main__":
    main()

import torch
from torch.utils.data import Dataset
import json
from minigpt.tokenizer import Tokenizer

class QADataset(Dataset):
    """
    自回归训练用的问答数据集（Question-Answer Dataset）。

    将 (question, answer) 拼接为单个序列，并构造：
      - 完整序列: [q1, q2, ..., <sep>, a1, a2, ..., <sep>, <pad>, ...]
      - input_ids: 完整序列去掉最后一个 token
      - targets:   完整序列去掉第一个 token
    """

    def __init__(self, data_path: str, tokenizer: Tokenizer, max_length: int):
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.data = []

        print(f"Loading data from {data_path}")

        with open(data_path, "r", encoding="utf-8") as f:
            for line_num, line in enumerate(f, 1):
                line = line.strip()
                if not line:
                    continue

                try:
                    item = json.loads(line)
                    if "question" not in item or "answer" not in item:
                        print(
                            f"⚠️  Line {line_num}: Missing 'question' or 'answer', skipped."
                        )
                        continue
                    self.data.append((item["question"], item["answer"]))

                except Exception as e:
                    print(f"⚠️  Line {line_num}: Invalid JSON, skipped. Error: {e}")

        print(f"✅ Loaded {len(self.data)} valid QA pairs.")

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        question, answer = self.data[idx]
        full_tokens, attn_mask = self.tokenizer.encode(
            question, answer, max_length=self.max_length
        )

        # 自回归训练： input 向右移一位，target 向左移一位
        input_ids = full_tokens[:-1]
        attention_mask = attn_mask[:-1]
        targets = full_tokens[1:]

        return {
            "input_ids": torch.tensor(input_ids, dtype=torch.long),
            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),
            "targets": torch.tensor(targets, dtype=torch.long),
        }

import torch
import torch.nn as nn
import numpy as np

class ScaledDotProductAttention(nn.Module):
    def __init__(self, d_k):
        super(ScaledDotProductAttention, self).__init__()
        self.d_k = d_k  # 输入的维度，用于计算缩放因子

    def forward(self, q, k, v, attn_mask):
        """
        Args:
            q: [batch_size, n_heads, len_q, d_k]
            k: [batch_size, n_heads, len_k, d_k]
            v: [batch_size, n_heads, len_v, d_v]  (注意: len_k == len_v)
            attn_mask: [batch_size, n_heads, len_q, len_k] 或 None
        Returns:
            context: [batch_size, n_heads, len_q, d_v]
            attn:    [batch_size, n_heads, len_q, len_k]
        """
        # step 1: 计算 QK^T / sqrt(d_k) shape: [batch_size, n_heads, len_q, len_k]
        scores = torch.matmul(q, k.transpose(-1, -2)) / np.sqrt(self.d_k)

        # step 2: 应用 attn_mask
        if attn_mask is not None:
            # masked_fill_: 将 mask=False 的位置替换为 -1e9
            scores = scores.masked_fill(~attn_mask, -1e9)

        # step 3: softmax 得到 attn
        attn = nn.Softmax(dim=-1)(scores)  # shape: [batch_size, n_heads, len_q, len_k]

        # step 4: attn * v
        context = torch.matmul(attn, v)  # shape: [batch_size, n_heads, len_q, head_dim]

        return context, attn

if __name__ == "__main__":
    # 用户偏好 Query
    q = torch.tensor([[[[0.9, 0.8]]]])

    # 三家饭店的 Key: [麻辣程度, 便宜程度]
    k = torch.tensor(
        [[[[0.1, 0.3], [0.9, 0.5], [0.3, 0.1]]]]
    )

    # 三家饭店的 Value: [评分, 出餐速度]
    v = torch.tensor(
        [[[[0.5, 0.9], [0.9, 0.6], [0.7, 0.8]]]]
    )

    mask = None

    attention = ScaledDotProductAttention(d_k=2)

    context, attn_weights = attention(q, k, v, mask)

    print("=== Attention Weights (α) ===")
    print(attn_weights.squeeze())

    print("\n=== Output Context (Weighted Value) ===")
    print(context.squeeze())

import torch
import torch.nn as nn
import numpy as np

from minigpt.attention_manual import ScaledDotProductAttention

class MultiHeadAttention(nn.Module):
    def __init__(self, n_heads, d_model, d_k, d_v):
        super(MultiHeadAttention, self).__init__()
        self.n_heads = n_heads
        self.d_model = d_model
        self.d_k = d_k
        self.d_v = d_v

        # 线性变换层
        self.w_q = nn.Linear(d_model, n_heads * d_k, bias=False)
        self.w_k = nn.Linear(d_model, n_heads * d_k, bias=False)
        self.w_v = nn.Linear(d_model, n_heads * d_v, bias=False)

        # 输出层
        self.fc = nn.Linear(n_heads * d_v, d_model, bias=False)

        # 注意力模块
        self.attention = ScaledDotProductAttention(d_k)

        # LayerNorm层
        self.layer_norm = nn.LayerNorm(d_model)

    def forward(self, q, k, v, attn_mask=None):
        """
        Args:
            q: [batch_size, len_q, d_model]
            k: [batch_size, len_k, d_model]
            v: [batch_size, len_v, d_model]  (len_k == len_v)
            attention_mask: [batch_size, len_q, len_k] or None
        Returns:
            output: [batch_size, len_q, d_model]
            attn:   [batch_size, n_heads, len_q, len_k]
        """
        residual = q
        batch_size = q.size(0)

        q_proj = self.w_q(q).view(batch_size, -1, self.n_heads, self.d_k)
        k_proj = self.w_k(k).view(batch_size, -1, self.n_heads, self.d_k)
        v_proj = self.w_v(v).view(batch_size, -1, self.n_heads, self.d_v)

        q_proj = q_proj.transpose(1, 2)
        k_proj = k_proj.transpose(1, 2)
        v_proj = v_proj.transpose(1, 2)

        if attn_mask is not None:
            attn_mask = attn_mask.unsqueeze(1).expand(-1, self.n_heads, -1, -1)
        else:
            attn_mask = None

        context, attn = self.attention(q_proj, k_proj, v_proj, attn_mask)

        context = (
            context.transpose(1, 2)
            .contiguous()
            .view(batch_size, -1, self.n_heads * self.d_v)
        )

        output = self.fc(context)
        output = self.layer_norm(output + residual)

        return output, attn

if __name__ == "__main__":
    torch.manual_seed(42)   # 固定种子，让输出可复现

    k_v_input = torch.tensor(
        [
            [
                [0.1, 0.3, 0.8, 0.5],  # 粤菜馆
                [0.9, 0.5, 0.2, 0.9],  # 川菜馆
                [0.3, 0.1, 0.9, 0.7],  # 日式餐厅
            ]
        ]
    )

    q_input = torch.tensor([[[0.9, 0.8, 0.3, 0.6]]])

    # 测试1: 无 mask
    mha = MultiHeadAttention(n_heads=2, d_model=4, d_k=2, d_v=2)
    output, attn_weights = mha(q_input, k_v_input, k_v_input, attn_mask=None)
    print("=== No Mask ===")
    print("Output shape:", output.shape)
    print("Attn shape:", attn_weights.shape)
    print("Head 1 weights:", attn_weights[0, 0, 0].tolist())
    print("Head 2 weights:", attn_weights[0, 1, 0].tolist())

    # 测试2: 屏蔽日式餐厅 (index=2)
    mask = torch.ones(1, 1, 3, dtype=torch.bool)
    mask[0, 0, 2] = False
    output2, attn_weights2 = mha(q_input, k_v_input, k_v_input, attn_mask=mask)
    print("\n=== With Mask ===")
    print("Head 1 weights:", attn_weights2[0, 0, 0].tolist())
    print("Head 2 weights:", attn_weights2[0, 1, 0].tolist())

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Optional, Union
from dataclasses import dataclass

@dataclass
class GPTConfig:
    """
    模型超参数配置类（使用 dataclass 简洁定义）

    这是一个 GPT 风格的缩小配置；隐藏维度与 GPT-2 Small 相同，
    但注意力头数和层数都更少。vocab_size 根据中文任务调整。
    """

    n_embd: int = 768  # token embedding 维度（也即 hidden size）
    n_head: int = 8    # 多头注意力头数（必须整除 n_embd）
    n_layer: int = 6   # Transformer block 层数
    dropout: float = 0.1  # 所有 dropout 层的丢弃率
    block_size: int = 256  # 最大上下文长度（位置编码最大支持长度）
    vocab_size: int = 4825  # 词表大小（根据实际 tokenizer 决定）

class CausalSelfAttention(nn.Module):
    def __init__(self, config: GPTConfig):
        super().__init__()

        assert config.n_embd % config.n_head == 0
        self.config = config
        self.head_dim = config.n_embd // config.n_head
        self.qkv_proj = nn.Linear(config.n_embd, 3 * config.n_embd)
        self.out_proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
        self.resid_dropout = nn.Dropout(config.dropout)

    def forward(
        self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
        B, T, C = x.shape

        qkv = self.qkv_proj(x)
        q, k, v = qkv.chunk(3, dim=-1)

        q = q.view(B, T, self.config.n_head, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.config.n_head, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.config.n_head, self.head_dim).transpose(1, 2)

        causal_mask = torch.tril(
            torch.ones(T, T, dtype=torch.bool, device=x.device)
        )[None, None, :, :]

        if attention_mask is not None:
            key_padding_mask = attention_mask.to(torch.bool)[:, None, None, :]
            attn_mask = causal_mask & key_padding_mask
        else:
            attn_mask = causal_mask

        y = F.scaled_dot_product_attention(
            q, k, v,
            attn_mask=attn_mask,
            dropout_p=self.config.dropout if self.training else 0,
            is_causal=False,
        )

        y = y.transpose(1, 2).contiguous().view(B, T, C)
        y = self.resid_dropout(self.out_proj(y))
        return y

import torch
from minigpt.qa_dataset import QADataset
from minigpt.tokenizer import Tokenizer
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
import torch.nn as nn
import os
from minigpt.model import GPTLMHeadModel, GPTConfig
from tqdm import tqdm

def train_epoch(model, train_loader, optimizer, criterion, device):
    model.train()
    total_loss = 0
    num_batches = len(train_loader)

    progress_bar = tqdm(
        enumerate(train_loader),
        total=num_batches,
        desc="  Train",
        leave=False,
        unit="batch",
    )

    for batch_idx, batch in progress_bar:
        input_ids = batch["input_ids"].to(device, non_blocking=True)
        attention_mask = batch["attention_mask"].to(device, non_blocking=True)
        targets = batch["targets"].to(device, non_blocking=True)

        optimizer.zero_grad()
        logits = model(input_ids=input_ids, attention_mask=attention_mask)
        loss = criterion(logits.view(-1, logits.size(-1)), targets.view(-1))

        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        progress_bar.set_postfix({"loss": f"{loss.item():.4f}"})

    return total_loss / num_batches

@torch.no_grad()
def validate(model, val_loader, criterion, device):
    model.eval()
    total_loss = 0
    num_batches = len(val_loader)

    progress_bar = tqdm(
        enumerate(val_loader),
        total=num_batches,
        desc="  Val",
        leave=False,
        unit="batch",
    )

    for batch_idx, batch in progress_bar:
        input_ids = batch["input_ids"].to(device, non_blocking=True)
        attention_mask = batch["attention_mask"].to(device, non_blocking=True)
        targets = batch["targets"].to(device, non_blocking=True)

        logits = model(input_ids=input_ids, attention_mask=attention_mask)
        loss = criterion(logits.view(-1, logits.size(-1)), targets.view(-1))
        total_loss += loss.item()
        progress_bar.set_postfix({"loss": f"{loss.item():.4f}"})

    return total_loss / num_batches

def save_checkpoint(model, optimizer, epoch, path):
    torch.save(
        {
            "epoch": epoch,
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
        },
        path,
    )

def train_model(
    model,
    train_loader,
    val_loader,
    optimizer,
    criterion,
    device,
    num_epochs,
    model_output_dir,
    writer,
):
    os.makedirs(model_output_dir, exist_ok=True)
    best_val_loss = float("inf")

    for epoch in range(1, num_epochs + 1):
        print(f"\nEpoch {epoch}/{num_epochs}")

        train_loss = train_epoch(model, train_loader, optimizer, criterion, device)
        print(f"  Train Loss: {train_loss:.4f}")

        val_loss = validate(model, val_loader, criterion, device)
        print(f"  Val Loss:   {val_loss:.4f}")

        if writer is not None:
            writer.add_scalar("Loss/train", train_loss, epoch)
            writer.add_scalar("Loss/val", val_loss, epoch)

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            best_path = os.path.join(model_output_dir, "best_model.pth")
            save_checkpoint(model, optimizer, epoch, best_path)
            print(f"  🎉 New best model saved (val loss: {val_loss:.4f})")

    print(f"\n✅ Training finished. Best validation loss: {best_val_loss:.4f}")

def main():
    train_path = "./data/train.jsonl"
    val_path = "./data/val.jsonl"
    vocab_path = "./data/vocab.json"

    max_length = 140
    batch_size = 64
    lr = 1e-4
    epochs = 15

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")
    if device.type == "cuda":
        print(f"GPU: {torch.cuda.get_device_name(0)}")

    tokenizer = Tokenizer(vocab_path)
    config = GPTConfig(vocab_size=tokenizer.get_vocab_size())
    model = GPTLMHeadModel(config).to(device)

    train_dataset = QADataset(train_path, tokenizer, max_length)
    val_dataset = QADataset(val_path, tokenizer, max_length)

    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=2,
        pin_memory=True,
        drop_last=True,
    )

    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=2,
        pin_memory=True,
        drop_last=False,
    )

    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.pad_token_id)

    writer = SummaryWriter("runs/minigpt")

    train_model(
        model=model,
        train_loader=train_loader,
        val_loader=val_loader,
        optimizer=optimizer,
        criterion=criterion,
        device=device,
        num_epochs=epochs,
        model_output_dir="output",
        writer=writer,
    )

    writer.close()
    print("\n🎉 Training pipeline completed.")

if __name__ == "__main__":
    main()

def generate(
    self,
    input_ids: torch.Tensor,
    max_new_tokens: int = 20,
    stop_token_ids: Optional[Union[int, List[int]]] = None,
    temperature: float = 1.0,
    top_k: Optional[int] = None,
    do_sample: bool = False,
) -> torch.Tensor:
    """
    自回归文本生成方法（支持多种解码策略）
    """
    self.eval()
    device = input_ids.device
    B, T = input_ids.shape

    stop_tokens = set()
    if stop_token_ids is not None:
        if isinstance(stop_token_ids, int):
            stop_tokens.add(stop_token_ids)
        else:
            stop_tokens.update(stop_token_ids)

    stop_tensor = None
    if stop_tokens:
        stop_tensor = torch.tensor(list(stop_tokens), device=device, dtype=input_ids.dtype)

    generated = input_ids.clone()
    finished = torch.zeros(B, dtype=torch.bool, device=device)

    with torch.no_grad():
        for _ in range(max_new_tokens):
            if finished.all() or generated.size(1) >= self.config.block_size:
                break

            logits = self(generated)
            next_token_logits = logits[:, -1, :]

            if do_sample:
                if temperature <= 0:
                    raise ValueError("temperature must be > 0")
                next_token_logits = next_token_logits / temperature

                if top_k is not None and top_k > 0:
                    k = min(top_k, next_token_logits.size(-1))
                    values, _ = torch.topk(next_token_logits, k, dim=-1)
                    threshold = values[:, -1:]
                    next_token_logits = torch.where(
                        next_token_logits < threshold,
                        torch.full_like(next_token_logits, float('-inf')),
                        next_token_logits
                    )

                probs = F.softmax(next_token_logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1)
            else:
                next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)

            last_token = generated[:, -1:].clone()
            next_token = torch.where(finished.unsqueeze(1), last_token, next_token)
            generated = torch.cat([generated, next_token], dim=1)

            if stop_tensor is not None:
                is_stop = (next_token == stop_tensor).any(dim=1)
                finished = finished | is_stop

    return generated

import torch
from minigpt.tokenizer import Tokenizer
from minigpt.model import GPTLMHeadModel, GPTConfig

def generate(model, tokenizer, question, max_new_tokens, device):
    """外部问答生成函数（使用模型内部的 generate 方法）"""
    tokens, _ = tokenizer.encode(
        question, answer=None, pad_to_max_length=False
    )
    input_ids = torch.tensor(tokens, dtype=torch.long, device=device).unsqueeze(0)

    sep_id = tokenizer.sep_token_id
    pad_id = tokenizer.pad_token_id

    output_ids = model.generate(
        input_ids=input_ids,
        max_new_tokens=max_new_tokens,
        stop_token_ids=[sep_id, pad_id],
        do_sample=False,
    )

    generated_ids = output_ids[0].tolist()

    try:
        first_sep = generated_ids.index(sep_id)
        answer_ids = generated_ids[first_sep + 1:]
    except ValueError:
        answer_ids = generated_ids[len(tokens):]

    answer_clean = []
    for tok in answer_ids:
        if tok != sep_id and tok != pad_id:
            answer_clean.append(tok)
        else:
            break

    return tokenizer.decode(answer_clean)

def main():
    vocab_path = "./data/vocab.json"
    max_new_tokens = 140
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model_path = "output/best_model.pth"

    tokenizer = Tokenizer(vocab_path)
    config = GPTConfig(vocab_size=tokenizer.get_vocab_size())
    model = GPTLMHeadModel(config).to(device)
    checkpoint = torch.load(model_path, map_location=device)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.eval()

    while True:
        question = input("Q: ")
        if question == "":
            break
        answer = generate(model, tokenizer, question, max_new_tokens, device)
        print("A:", answer)

if __name__ == "__main__":
    main()

if T > block_size:
    raise ValueError(
        f"prompt 长度 {T} 超过 block_size {block_size}；"
        f"请截断 prompt，或在 GPTConfig 中调大 block_size"
    )

if answer is not None:
    body = [self.word2id.get(c, self.unk_token_id) for c in answer]
    if keep_eos:
        budget = max_length - len(tokens) - 1   # 给 <sep> 留一个位置
        if budget < 0:
            tokens = tokens[: max_length - 1]
            budget = 0
        tokens.extend(body[:budget])
        tokens.append(self.sep_token_id)
    else:
        tokens.extend(body)
        tokens.append(self.sep_token_id)

def decode(self, ids, skip_special: bool = True) -> str:
    out = []
    for i in ids:
        token = self.id2word[int(i)]
        if token == self.id2word[self.pad_token_id]:
            continue
        if skip_special and token in ("<pad>", "<unk>", "<sep>"):
            continue
        out.append(token)
    return "".join(out)