Update app.py

app.py

@@ -2,14 +2,35 @@ import torch
 import json
 import gradio as gr
 
+TITLE = "Danboru Tag Similarity"
+DESCRIPTION = """
+与えられたダンボールタグの類似度を計算します。\n
+対応するタグのリストはFilesからそれぞれのテキストファイルを参照してください。(Dartと同じです）。\n
+Dartを参考に、isek-ai/danbooru-tags-2023データセットでタグをシャッフルして2エポック学習しました。\n
+学習後のトークン埋め込みを元に計算しています。
+"""
+
 with open("num_to_token.json", "r") as f:
     num_to_token = json.load(f)
+with open("popular.txt", "r") as f:
+    populars = f.read().splitlines()
+with open("character.txt", "r") as f:
+    characters = f.read().splitlines()
+    characters_populars = list(set(characters) & set(populars))
+with open("copyright.txt", "r") as f:
+    copyrights = f.read().splitlines()
+    copyrights_populars = list(set(copyrights) & set(populars))
+with open("general.txt", "r") as f:
+    generals = f.read().splitlines()
+    generals_populars = list(set(generals) & set(populars))
+
+    
 token_to_num = {v:k for k,v in num_to_token.items()}
 token_embeddings = torch.load("token_embeddings.pt")
 
 tags = sorted(list(num_to_token.values()))
 
-def predict(target_tag, sort_by="descend"):
+def predict(target_tag, sort_by, category, popular):
     if sort_by == "descending":
         multiplier = 1
     else:
@@ -17,16 +38,28 @@ def predict(target_tag, sort_by="descend"):
     target_embedding = token_embeddings[int(token_to_num[target_tag])].unsqueeze(0)
     sims = torch.cosine_similarity(target_embedding, token_embeddings, dim=1)
     results = {num_to_token[str(i)]:sims[i].item() * multiplier for i in range(len(num_to_token))}
-
-    return results
+    if category == "general":
+        tag_list = generals if popular == "all" else generals_populars
+    elif category == "character":
+        tag_list = characters if popular == "all" else characters_populars
+    elif category == "copyright":
+        tag_list = copyrights if popular == "all" else copyrights_populars
+    else:
+        tag_list = results.keys() if popular == "all" else populars
+        
+    return {k:results[k] for k in tag_list}
 
 demo = gr.Interface(
         fn=predict, 
         inputs=[
             gr.Textbox(label="Target tag", value="otoko no ko"),
-            gr.Radio(choices=["descending", "ascending"], label="Sort by", value="descending")
+            gr.Radio(choices=["descending", "ascending"], label="Sort by", value="descending"),
+            gr.Dropdown(choices=["all", "general", "character", "copyright"], value="all", label="category"),
+            gr.Radio(choices=["all", "only_popular"], label="Only popular tag (count>=1000)", value="all"),
         ],
         outputs=gr.Label(num_top_classes=50),
+        title=TITLE,
+        description=DESCRIPTION
     )
 
-demo.launch()
+demo.launch()