Create app.py

app.py

@@ -0,0 +1,32 @@
+import torch
+import json
+import gradio as gr
+
+with open("num_to_token.json", "r") as f:
+    num_to_token = json.load(f)
+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"):
+    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))}
+
+    return results
+
+demo = gr.Interface(
+        fn=predict, 
+        inputs=[
+            gr.Dropdown(choices=tags, label="Target tag", value="otoko no ko"),
+            gr.Dropdown(choices=["ascending", "descending"], label="Sort by", value="descending")
+        ],
+        outputs=gr.Label(num_top_classes=50),
+    )
+
+demo.launch()