Update app.py

app.py

@@ -10,44 +10,42 @@ Dartを参考に、isek-ai/danbooru-tags-2023データセットでタグをシ
 学習後のトークン埋め込みを元に計算しています。
 """
 
-with open("num_to_token.json", "r") as f:
-    num_to_token = json.load(f)
+with open("id_to_token.json", "r") as f:
+    id_to_token = json.load(f)
+token_to_id = {v:int(k) for k,v in id_to_token.items()}
+
 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()}
+tags = characters + copyrights + generals
 token_embeddings = torch.load("token_embeddings.pt")
 
-tags = sorted(list(num_to_token.values()))
-
 def predict(target_tag, sort_by, category, popular):
     if sort_by == "descending":
         multiplier = 1
     else:
         multiplier = -1
-    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))}
+    target_embedding = token_embeddings[token_to_id[target_tag]].unsqueeze(0)
+    sims = torch.cosine_similarity(target_embedding, token_embeddings, dim=1) * multiplier
+    
     if category == "general":
-        tag_list = generals if popular == "all" else generals_populars
+        tag_list = generals
     elif category == "character":
-        tag_list = characters if popular == "all" else characters_populars
+        tag_list = characters
     elif category == "copyright":
-        tag_list = copyrights if popular == "all" else copyrights_populars
+        tag_list = copyrights
     else:
-        tag_list = results.keys() if popular == "all" else populars
+        tag_list = tags
+    
+    if popular=="only_popular":
+        tag_list = list(set(tag_list) & set(populars))
         
-    return {k:results[k] for k in tag_list}
+    return {k:sims[token_to_id[k]].item() for k in tag_list}
 
 demo = gr.Interface(
         fn=predict,