Modified model loading
Browse files- app/model.py +2 -2
app/model.py
CHANGED
@@ -3,7 +3,7 @@ import os
|
|
3 |
import re
|
4 |
|
5 |
class CancerClassifier:
|
6 |
-
def __init__(self, model_path
|
7 |
model = AutoModelForSequenceClassification.from_pretrained(model_path)
|
8 |
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
9 |
self.classifier = pipeline(
|
@@ -25,7 +25,7 @@ class CancerClassifier:
|
|
25 |
}
|
26 |
|
27 |
class CancerExtractor:
|
28 |
-
def __init__(self, model_path
|
29 |
self.extractor = pipeline(
|
30 |
"ner",
|
31 |
model=model_path,
|
|
|
3 |
import re
|
4 |
|
5 |
class CancerClassifier:
|
6 |
+
def __init__(self, model_path="user1729/BiomedBERT-cancer-bert-classifier-v1.0"):
|
7 |
model = AutoModelForSequenceClassification.from_pretrained(model_path)
|
8 |
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
9 |
self.classifier = pipeline(
|
|
|
25 |
}
|
26 |
|
27 |
class CancerExtractor:
|
28 |
+
def __init__(self, model_path="alvaroalon2/biobert_diseases_ner"):
|
29 |
self.extractor = pipeline(
|
30 |
"ner",
|
31 |
model=model_path,
|