Merge branch 'main' of https://huggingface.co/michael-0acf4/anitag2vec

README.md

@@ -1,44 +1,46 @@
----
-license: mit
----
-
-# AniTag2Vec
-
-Training and inference examples are all available on [my github](https://github.com/michael-0acf4/anitag2vec).
-
-Implementation is detailed in [this blog post](https://blog.afmichael.dev/posts/2026/set-embeddings-and-anitag2vec/).
-
-```python
-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'
-```
+---
+license: mit
+---
+
+# AniTag2Vec
+
+Generate vector embeddings from Danbooru, Sakugabooru, Pixiv, MAL style tags.
+
+Training and inference examples are all available on [my github](https://github.com/michael-0acf4/anitag2vec).
+
+Implementation is detailed in [this blog post](https://blog.afmichael.dev/posts/2026/set-embeddings-and-anitag2vec/).
+
+```python
+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'
+```