Spaces:
Sleeping
Sleeping
Update load_model.py
Browse files- load_model.py +2 -2
load_model.py
CHANGED
@@ -14,7 +14,7 @@ def load_model(model_name):
|
|
14 |
return model
|
15 |
|
16 |
def load_tokenizer(model_name):
|
17 |
-
tokenizer = BertTokenizerFast.from_pretrained(model_name, use_auth_token=os.getenv('API_KEY')
|
18 |
return tokenizer
|
19 |
|
20 |
def predict(text, model, tokenizer):
|
@@ -23,7 +23,7 @@ def predict(text, model, tokenizer):
|
|
23 |
return outputs
|
24 |
|
25 |
def main():
|
26 |
-
model_name = os.getenv('
|
27 |
if model_name is None:
|
28 |
raise ValueError("MODEL_PATH environment variable not set or is None")
|
29 |
|
|
|
14 |
return model
|
15 |
|
16 |
def load_tokenizer(model_name):
|
17 |
+
tokenizer = BertTokenizerFast.from_pretrained(model_name, use_auth_token=os.getenv('API_KEY'))
|
18 |
return tokenizer
|
19 |
|
20 |
def predict(text, model, tokenizer):
|
|
|
23 |
return outputs
|
24 |
|
25 |
def main():
|
26 |
+
model_name = os.getenv('MODEL_PATH')
|
27 |
if model_name is None:
|
28 |
raise ValueError("MODEL_PATH environment variable not set or is None")
|
29 |
|