pipeline update

README.md

@@ -6,7 +6,7 @@ colorTo: purple
 sdk: streamlit
 sdk_version: 1.2.0
 app_file: app.py
-models: ["dslim/bert-base-NER"]
+models: ["Jean-Baptiste/roberta-large-ner-english"]
 pinned: false
 license: mit
 ---

app.py

@@ -27,8 +27,8 @@ COLOR = {
 @st.cache(allow_output_mutation=True, show_spinner=True)
 def load_models():
     # NER
-    tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
-    bert_ner = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER")
+    tokenizer = AutoTokenizer.from_pretrained("Jean-Baptiste/roberta-large-ner-english")
+    bert_ner = AutoModelForTokenClassification.from_pretrained("Jean-Baptiste/roberta-large-ner-english")
     tagger = pipeline("token-classification", model=bert_ner, tokenizer=tokenizer,
                       device=-1, aggregation_strategy="average")
     # NED
@@ -54,8 +54,8 @@ def get_candidates(mentions_tags):
         if (mention, tag) in cache.keys():
             candidates.append((mention, cache[(mention, tag)]))
         else:
-            res1 = google_search(mention + TYPE[tag], limit=1)
-            res2 = wikipedia_search(mention, limit=5)
+            res1 = google_search(mention + TYPE[tag], limit=3)
+            res2 = wikipedia_search(mention, limit=3)
             cands = list(set(res1 + res2))
             cache[(mention, tag)] = cands
             candidates.append((mention, cands))