test.py

import torch
import torch.nn as nn
from model.LLM import SimpleGPT, visualize_tokens_line
from tokenizers import Tokenizer


if __name__ == '__main__':
    max_length = 256
    step = 1
    vocab_size = 15099
    output_dim = 256
    tokenizer = Tokenizer.from_file("./Tokenizer/tokenizer.json")
    torch.manual_seed(1234)
    PATH = "./GPT.pth"

    text = "Hello, how are "
    token_ids = torch.tensor(tokenizer.encode(text).ids).unsqueeze(0).cuda()
    print(visualize_tokens_line(token_ids, tokenizer, "test"))

    model = SimpleGPT(vocab_size, max_length, output_dim)
    model.load_state_dict(torch.load(PATH))
    model.eval()
    model.cuda()
    with torch.no_grad():
        for _ in range(50):  # 生成50个token
            out = model(token_ids)
            out = out[:, -1, :]
            probas = torch.softmax(out, dim=-1)
            idx_next = torch.argmax(probas, dim=-1, keepdim=True)

            token_ids = torch.cat((token_ids, idx_next), dim=1)

    print(tokenizer.decode(token_ids[0].tolist()))