Spaces:
Runtime error
Runtime error
AhdCompnay
commited on
Commit
•
84b31d5
1
Parent(s):
07cbdb5
Update kpe_ranker.py
Browse files- kpe_ranker.py +6 -2
kpe_ranker.py
CHANGED
@@ -7,10 +7,14 @@ from huggingface_hub import hf_hub_download
|
|
7 |
|
8 |
class KpeRanker:
|
9 |
def __init__(self):
|
|
|
|
|
|
|
|
|
10 |
local_dir = "./"
|
11 |
-
model_path = os.path.join(local_dir,
|
12 |
if not os.path.isfile(model_path):
|
13 |
-
hf_hub_download(repo_id=
|
14 |
TRAINED_MODEL_ADDR = model_path
|
15 |
# TRAINED_MODEL_ADDR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'trained_model', 'trained_model_10000.pt')
|
16 |
self.kpe = KPE(trained_kpe_model= TRAINED_MODEL_ADDR, flair_ner_model='flair/ner-english-ontonotes-large', device='cpu')
|
|
|
7 |
|
8 |
class KpeRanker:
|
9 |
def __init__(self):
|
10 |
+
model_name = os.environ.get("MODEL_NAME")
|
11 |
+
model_repo = os.environ.get("MODEL_REPO")
|
12 |
+
model_token = os.environ.get("MODEL_TOKEN")
|
13 |
+
|
14 |
local_dir = "./"
|
15 |
+
model_path = os.path.join(local_dir, model_name)
|
16 |
if not os.path.isfile(model_path):
|
17 |
+
hf_hub_download(repo_id=model_repo, filename=model_name, local_dir=local_dir, token=model_token)
|
18 |
TRAINED_MODEL_ADDR = model_path
|
19 |
# TRAINED_MODEL_ADDR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'trained_model', 'trained_model_10000.pt')
|
20 |
self.kpe = KPE(trained_kpe_model= TRAINED_MODEL_ADDR, flair_ner_model='flair/ner-english-ontonotes-large', device='cpu')
|