jsylee commited on
Commit
2600008
1 Parent(s): 850577a
Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -10,10 +10,12 @@ model_checkpoint = "jsylee/scibert_scivocab_uncased-finetuned-ner"
10
  model = AutoModelForTokenClassification.from_pretrained(model_checkpoint,
11
  num_labels=5,
12
  id2label={0: 'O', 1: 'DRUG', 2: 'DRUG', 3: 'ADVERSE EFFECT', 4: 'ADVERSE EFFECT'} # for grouping BIO tags back together
13
- ).to("cpu")
14
  tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
15
 
16
- model_pipeline = pipeline(task="ner", model=model, tokenizer=tokenizer, device=0, grouped_entities=True)
 
 
17
 
18
  def extract_entities(sentence):
19
  """ Extract drug and reaction entities, and show using displaCy's NER visualizer.
 
10
  model = AutoModelForTokenClassification.from_pretrained(model_checkpoint,
11
  num_labels=5,
12
  id2label={0: 'O', 1: 'DRUG', 2: 'DRUG', 3: 'ADVERSE EFFECT', 4: 'ADVERSE EFFECT'} # for grouping BIO tags back together
13
+ )
14
  tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
15
 
16
+ model.to("cpu")
17
+
18
+ model_pipeline = pipeline(task="ner", model=model, tokenizer=tokenizer, device=-1, grouped_entities=True)
19
 
20
  def extract_entities(sentence):
21
  """ Extract drug and reaction entities, and show using displaCy's NER visualizer.