fixed torch device

README.md

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

app.py

@@ -5,7 +5,6 @@ from transformers import AutoModelForTokenClassification, AutoTokenizer, pipelin
 
 from src import GBRT, wikipedia_search, google_search
 
-
 TYPE = {
     'LOC': ' location',
     'PER': ' person',
@@ -31,7 +30,7 @@ def load_models():
     tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
     bert_ner = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER")
     tagger = pipeline("token-classification", model=bert_ner, tokenizer=tokenizer,
-                      device=0, aggregation_strategy="average")
+                      device=-1, aggregation_strategy="average")
     # NED
     model = GBRT()
     return model, tagger

requirements.txt

@@ -2,5 +2,6 @@ nltk
 numpy
 pandas
 wikipedia2vec
-torch
-transformers
+transformers
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch