Update app.py

app.py

@@ -19,6 +19,12 @@ general_tags = df[df["category"] == 0].index
 character_tags = df[df["category"] == 4].index
 all_tags = df.index
 
+tag_pair_df = pd.read_parquet("hf://datasets/p1atdev/danbooru-ja-tag-pair-20241015/data/train-00000-of-00001.parquet")
+tag_pair = {title:other_names[0] for title, other_names in zip(tag_pair_df["title"], tag_pair_df["other_names"])}
+for tag in df["name"]:
+    if tag not in tag_pair:
+        tag_pair[tag] = ""
+
 def predict(target_tags, search_in):
     target_tags = [tag.strip().replace(" ", "_") for tag in target_tags.split(",")]
     target_ids = [label2id[tag] for tag in target_tags]
@@ -26,7 +32,7 @@ def predict(target_tags, search_in):
 
     sim = torch.cosine_similarity(query, head.unsqueeze(0), dim=2).mean(dim=0)
     tags = general_tags if search_in == "general" else character_tags if search_in == "character" else all_tags
-    return {id2label[i]: sim[i].item() for i in tags}
+    return {f"{id2label[i]}({tag_pair[id2label[i]]})": sim[i].item() for i in tags}
 
 demo = gr.Interface(
         fn=predict,