Spaces:
Runtime error
Runtime error
typo
Browse files
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 |
-
)
|
14 |
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
|
15 |
|
16 |
-
|
|
|
|
|
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.
|