metadata
license: mit
AniTag2Vec
Training and inference examples are all available on my github.
Implementation is detailed in this blog post.
TOKENIZER_PATH = "./checkpoints/token_dataset_c7359727bcee4f8b_vocab_size_5000_freq_3.json"
CONFIG_PATH = "./checkpoints/setup_params_8ea07c7d34b64b69_c7359727bcee4f8b.json"
MODEL_PATH = "./checkpoints/anitag2vec_8ea07c7d34b64b69_c7359727bcee4f8b_i128_e20_s60203_b256_p1871744.pth"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
cfg = SetupConfig.load_from_file(CONFIG_PATH)
print(cfg)
tagtok = TagBPETokenizer(vocab_size=cfg.HYPERP_TAGTOK_VOCAB_SIZE, min_frequency=cfg.HYPERP_TAGTOK_MIN_FREQ)
tagtok.load(TOKENIZER_PATH)
anitag2vec = AniTag2Vec(
vocab_size=tagtok.vocab_size,
max_len_cut=cfg.HYPERP_TAGTOK_MAX_TOKEN_CLAMP,
d_model=cfg.HYPERP_TRANSFORMER_D_MODEL,
n_heads=cfg.HYPERP_TRANSFORMER_N_HEADS,
n_layers=cfg.HYPERP_TRANSFORMER_N_LAYERS,
output_emb=cfg.HYPERP_OUTPUT_EMB,
)
anitag2vec.to(device)
anitag2vec.load_state_dict(torch.load(MODEL_PATH))
anitag2vec.eval()
runner = AniTag2VecRunner(tagtok, anitag2vec)
# Inference
def compare(a: str, b: str):
ax = runner.run_inference_human([a])
bx = runner.run_inference_human([b])
howmuch = ((F.normalize(ax) @ F.normalize(bx).T).item())
print(f"{howmuch:.2f}: '{a}' vs '{b}'")
compare("#1girl #1boy", "#1boy #1girl")
# 1.00: '#1girl #1boy' vs '#1boy #1girl'