Filter meaningless images

Author: Linyxus

app.py

@@ -16,4 +16,20 @@
 @app.get('/search/')
 def search(q: str):
     results = engine.search(q)
-    return [x.document.main_image for x in results]
+    res = [x.document.main_image for x in results]
+
+    def is_bullshit(x: str) -> bool:
+        if x.endswith('20px-Semi-protection-shackle.svg.png'):
+            return True
+        if x.endswith('50px-Question_book-new.svg.png'):
+            return True
+        if x.endswith('40px-Edit-clear.svg.png'):
+            return True
+        if x.endswith('50px-Question_book-new.svg.png'):
+            return True
+        if x.endswith('19px-Symbol_support_vote.svg.png'):
+            return True
+        return False
+
+    res = [x for x in res if not is_bullshit(x)]
+    return res

wiki_search/dataset/dataset.py

@@ -7,6 +7,8 @@
 import os
 from nltk.corpus import stopwords
 from nltk.tokenize import word_tokenize
+from torch_geometric.utils import degree
+from torch_scatter import scatter
 import torch
 import ray
 
@@ -79,6 +81,20 @@ def derive_edges_runner(doc: Document, d2i: Dict[str, int]):
     return derive_edges(doc, d2i)
 
 
+def compute_pr(edge_index, damp: float = 0.85, k: int = 10):
+    num_nodes = edge_index.max().item() + 1
+    deg_out = degree(edge_index[0])
+    x = torch.ones((num_nodes, )).to(edge_index.device).to(torch.float32)
+
+    for i in range(k):
+        edge_msg = x[edge_index[0]] / deg_out[edge_index[0]]
+        agg_msg = scatter(edge_msg, edge_index[1], reduce='sum')
+
+        x = (1 - damp) * x + damp * agg_msg
+
+    return x
+
+
 def ray_reduce(func, xs):
     n = len(xs)
     if n == 1: