jiangjiechen commited on
Commit
964af23
1 Parent(s): 20e9c0d

fix device issue: cpu

Browse files
src/parsing_client/sentence_parser.py CHANGED
@@ -6,6 +6,7 @@
6
 
7
  import re
8
  import nltk
 
9
  from nltk.stem import WordNetLemmatizer
10
  from allennlp.predictors.predictor import Predictor
11
 
@@ -113,7 +114,7 @@ def refine_results(tokens, spans, stopwords):
113
 
114
 
115
  class SentenceParser:
116
- def __init__(self, device='cuda:0',
117
  ner_path="https://storage.googleapis.com/allennlp-public-models/ner-model-2020.02.10.tar.gz",
118
  cp_path="https://storage.googleapis.com/allennlp-public-models/elmo-constituency-parser-2020.02.10.tar.gz"):
119
  self.device = self.parse_device(device)
 
6
 
7
  import re
8
  import nltk
9
+ import torch.cuda
10
  from nltk.stem import WordNetLemmatizer
11
  from allennlp.predictors.predictor import Predictor
12
 
 
114
 
115
 
116
  class SentenceParser:
117
+ def __init__(self, device='cuda' if torch.cuda.is_available() else 'cpu',
118
  ner_path="https://storage.googleapis.com/allennlp-public-models/ner-model-2020.02.10.tar.gz",
119
  cp_path="https://storage.googleapis.com/allennlp-public-models/elmo-constituency-parser-2020.02.10.tar.gz"):
120
  self.device = self.parse_device(device)