Create app.py

app.py

@@ -0,0 +1,42 @@
+import gradio as gr
+import timm
+import numpy as np
+import faiss
+import pandas as pd
+
+
+TITLE = "wd-eva02-large-tagger-v3-vector"
+DESCRIPTION = """
+[model](https://huggingface.co/SmilingWolf/wd-eva02-large-tagger-v3)
+"""
+
+model = timm.create_model(f"hf_hub:SmilingWolf/wd-eva02-large-tagger-v3", pretrained=True)
+head = model.head.weight.data.cpu().numpy()
+del model
+df = pd.read_csv(f"https://huggingface.co/SmilingWolf/wd-eva02-large-tagger-v3/resolve/main/selected_tags.csv")
+id2label = df["name"].to_dict()
+label2id = {v:k for k,v in id2label.items()}
+
+faiss.normalize_L2(head)
+index = faiss.IndexFlatIP(head.shape[1])
+index.add(head)
+
+def predict(target_tag):
+    target_id = label2id[target_tag]
+    query = head[target_id:target_id+1]
+    k = 50
+    target_id = label2id[target_tag]
+    distances, indices = index.search(query, k)
+    return {id2label[indice]:distance for indice, distance in zip(indices[0], distances[0])}
+
+demo = gr.Interface(
+        fn=predict, 
+        inputs=[
+            gr.Dropdown(list(label2id.keys()), label="Target tag", value="otoko_no_ko"),
+        ],
+        outputs=gr.Label(num_top_classes=50),
+        title=TITLE,
+        description=DESCRIPTION
+    )
+
+demo.launch()