Views
No views yet
1TOKENIZER_PATH = "./pytorch/token_dataset_b0d065e705028cb3_vocab_size_5000_freq_3.json"
2CONFIG_PATH = "./pytorch/config_63fc21b89723d1ce_b0d065e705028cb3.json"
3MODEL_PATH = "./pytorch/anitag2vec_63fc21b89723d1ce_b0d065e705028cb3_i128_e30_s157043_b256_p1871744.pth"
4
5device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
6cfg = ModelConfig.load_from_file(CONFIG_PATH)
7tagtok = TagBPETokenizer.load_from_file(TOKENIZER_PATH)
8
9anitag2vec = AniTag2Vec(
10 vocab_size=cfg.HYPERP_TAGTOK_VOCAB_SIZE,
11 max_len_cut=cfg.HYPERP_TAGTOK_MAX_TOKEN_CLAMP,
12 d_model=cfg.HYPERP_TRANSFORMER_D_MODEL,
13 n_heads=cfg.HYPERP_TRANSFORMER_N_HEADS,
14 n_layers=cfg.HYPERP_TRANSFORMER_N_LAYERS,
15 output_emb=cfg.HYPERP_OUTPUT_EMB,
16)
17anitag2vec.to(device)
18anitag2vec.load_state_dict(torch.load(MODEL_PATH))
19anitag2vec.eval()
20runner = AniTag2VecRunner(tagtok, anitag2vec)
21
22# Inference
23def compare(a: str, b: str):
24 ax = runner.run_inference_human([a])
25 bx = runner.run_inference_human([b])
26 howmuch = ((F.normalize(ax) @ F.normalize(bx).T).item())
27 print(f"{howmuch:.2f}: '{a}' vs '{b}'")
28
29compare("#1girl #1boy", "#1boy #1girl")
30# 1.00: '#1girl #1boy' vs '#1boy #1girl'