pipeline update

app.py

@@ -1,9 +1,8 @@
-from turtle import color
 import pandas as pd
 import streamlit as st
 from transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline
 
-from src import GBRT, wikipedia_search, google_search
+from src import GBRT, google_search, wikidata_search
 
 TYPE = {
     'LOC': ' location',
@@ -55,7 +54,7 @@ def get_candidates(mentions_tags):
             candidates.append((mention, cache[(mention, tag)]))
         else:
             res1 = google_search(mention + TYPE[tag], limit=3)
-            res2 = wikipedia_search(mention, limit=3)
+            res2 = wikidata_search(mention, limit=3)
             cands = list(set(res1 + res2))
             cache[(mention, tag)] = cands
             candidates.append((mention, cands))

src/__init__.py

@@ -60,23 +60,53 @@ def cosine_similarity(v1, v2):
         return np.dot(v2, v1) / v1v2
 
 
-def wikipedia_search(query, limit=20):
-    service_url = 'https://en.wikipedia.org/w/api.php'
+def is_disamb_page(title):
+    service_url = "https://en.wikipedia.org/w/api.php"
     params = {
-        'action': 'opensearch',
-        'search': query,
-        'namespace': 0,
-        'limit': limit,
-        'redirects': 'resolve',
+        "action": "query",
+        "prop": "pageprops",
+        "ppprop" : "disambiguation",
+        "redirects":'',
+        "format": "json",
+        "titles": title
+    }
+    results = requests.get(service_url, params=params).json()
+    return 'disambiguation' in str(results)
+
+
+def wikidata_search(query, limit=3):
+    service_url = 'https://www.wikidata.org/w/api.php'
+    params1 = {
+        "action": "wbsearchentities",
+        "search": query,
+        "language": "en",
+        "limit": limit,
+        "format": "json"
     }
 
-    results = requests.get(service_url, params=params).json()[1]
-    results = [i.replace(' ', '_')
-               for i in results if 'disambiguation' not in i.lower()]
-    return results
+    params2 = {
+        "action": "wbgetentities",
+        "language": "en",
+        "props": "sitelinks",
+        "sitefilter": "enwiki",
+        "format": "json"
+    }
+
+    results = requests.get(service_url, params=params1).json()
+    entities = [i['id'] for i in results['search']]
+
+    params2['ids'] = '|'.join(entities)
+    results = requests.get(service_url, params=params2).json()
+    candidates = []
+    for i in entities:
+        try:
+            candidates.append(results['entities'][i]['sitelinks']['enwiki']['title'].replace(' ', '_'))
+        except:
+            pass
+    return [i for i in candidates if is_disamb_page(i) == False]
 
 
-def google_search(query, limit=10):
+def google_search(query, limit=3):
     service_url = "https://www.googleapis.com/customsearch/v1/siterestrict"
     params = {
         'q': query,
@@ -88,6 +118,7 @@ def google_search(query, limit=10):
     res = requests.get(service_url, params=params)
     try:
         cands = [i['title'].replace(' - Wikipedia', '') for i in res.json()["items"]]
+        cands = [i for i in cands if is_disamb_page(i) == False]
         return [i.replace(' ', '_') for i in cands]
     except:
         return []