Witllm/wit/inference.py

import pytorch_lightning as pl
import torch

from model.qwen_module import QwenModule
from model.modeling_wit import QwenRunner
from model.tokenization_qwen import QWenTokenizer
import numpy as np

import configuration
import dataset.dataset as ds
import dataset.node_tree as nt

if __name__ == "__main__":

    checkpoint_path = "log/bigger/version_0/checkpoints/epoch=19-step=98720.ckpt"

    qwen = QwenModule.load_from_checkpoint(checkpoint_path=checkpoint_path)
    qwen.eval()
    conf = qwen.config
    torch.manual_seed(conf.seed)
    np.random.seed(conf.seed)
    runner = QwenRunner(qwen.llm)

    # batch = torch.tensor([[41]], dtype=torch.int64)
    # print(runner.ChatTokens(batch).detach().cpu().numpy()[0])

    val = ds.InitValDataset(conf).dataset
    md = val.meaning_dataset

    map = md.get_meaning_map()

    item = md.get_token(0)

    node = map.get_tree(md.get_meaning(0))
    # node.print()

    for i in range(1, len(item)):
        itemm = [item[:i]]
        batch = torch.tensor([item[:i]], dtype=torch.int64)
        next_token = runner.ChatTokens(batch, sample=False).detach().cpu().numpy()[0]
        if item[i] != next_token:
            node.set_seq_prop(i, "ERR_" + str(next_token))
            print(str(item[i]) + "  " + str(next_token) + "  ERROR")
    node.print()