Views
No views yet
lapvqa-pretrain-captioning).
Checkpoint format: {state_dict, vis_dim, d_model, num_layers, nhead, encoder, epoch, val_bleu4}.1import torch
2import tiktoken
3from lapvqa.rrg.heads import ReportGenerationHead
4
5ckpt = torch.load("pretrain-captioning.pt", map_location="cpu")
6head = ReportGenerationHead(
7 vis_dim = ckpt["vis_dim"],
8 d_model = ckpt["d_model"],
9 num_layers = ckpt["num_layers"],
10 nhead = ckpt["nhead"],
11)
12head.load_state_dict(ckpt["state_dict"])
13head.eval()
14
15enc = tiktoken.get_encoding("gpt2")
16bos_id = eos_id = enc.eot_token
17# pair with encoder_final.pt from lapvqa-pretrain-captioning
18token_ids = head.generate(vis_tokens, bos_id=bos_id, eos_id=eos_id)
19reports = [enc.decode(ids) for ids in token_ids]