Spaces:
Sleeping
Sleeping
Update load_model.py
Browse files- load_model.py +2 -2
load_model.py
CHANGED
@@ -8,13 +8,13 @@ load_dotenv()
|
|
8 |
def load_model(model_name):
|
9 |
try:
|
10 |
# Load TensorFlow model from Hugging Face
|
11 |
-
model = TFBertForSequenceClassification.from_pretrained(model_name, use_auth_token=os.getenv('
|
12 |
except OSError:
|
13 |
raise ValueError("Model loading failed.")
|
14 |
return model
|
15 |
|
16 |
def load_tokenizer(model_name):
|
17 |
-
tokenizer = BertTokenizerFast.from_pretrained(model_name, use_auth_token=os.getenv('
|
18 |
return tokenizer
|
19 |
|
20 |
def predict(text, model, tokenizer):
|
|
|
8 |
def load_model(model_name):
|
9 |
try:
|
10 |
# Load TensorFlow model from Hugging Face
|
11 |
+
model = TFBertForSequenceClassification.from_pretrained(model_name, use_auth_token=os.getenv('API_KEY'), from_tf=True)
|
12 |
except OSError:
|
13 |
raise ValueError("Model loading failed.")
|
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):
|