pipeline update

app.py

@@ -2,7 +2,7 @@ import pandas as pd
 import streamlit as st
 from transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline
 
-from src import GBRT, wikipedia_search, wikidata_search
+from src import GBRT, wikipedia_search, wikidata_search, google_search
 
 TYPE = {
     'LOC': ' location',
@@ -53,6 +53,7 @@ def get_candidates(mentions_tags):
             candidates.append((mention, cache[(mention, tag)]))
         else:
             cands = wikidata_search(mention, limit=3)
+            cands = list(set(cands + google_search(mention)))
             if cands == []:
                 cands = wikipedia_search(mention, limit=3)
             cache[(mention, tag)] = cands

src/__init__.py

@@ -121,6 +121,24 @@ def wikipedia_search(query, limit=3):
     return [i for i in results if is_disamb_page(i) == False]
 
 
+def google_search(query, limit=10):
+    service_url = "https://www.googleapis.com/customsearch/v1/siterestrict"
+    params = {
+        'q': query,
+        'num': limit,
+        'start': 0,
+        'key': os.environ.get('APIKEY'),
+        'cx': os.environ.get('CESCX')
+    }
+    res = requests.get(service_url, params=params)
+    try:
+        cands = [i['title'].replace(' - Wikipedia', '') for i in res.json()["items"]]
+        return [i.replace(' ', '_') for i in cands if is_disamb_page(i) == False]
+    except:
+        return []
+
+
+
 def get_entity_extract(entity_title, num_sentences=0):
     service_url = 'https://en.wikipedia.org/w/api.php'
     params = {