Spaces:
Running
Running
Update load_model.py
Browse files- load_model.py +19 -19
load_model.py
CHANGED
@@ -1,19 +1,19 @@
|
|
1 |
-
from transformers import BertTokenizer, BertForSequenceClassification
|
2 |
-
import streamlit as st
|
3 |
-
|
4 |
-
# Dictionary to map model names to their paths
|
5 |
-
model_paths = {
|
6 |
-
"cahya/bert-base-indonesian-522M": "
|
7 |
-
"indobenchmark/indobert-base-p2": "
|
8 |
-
"indolem/indobert-base-uncased": "
|
9 |
-
"mdhugol/indonesia-bert-sentiment-classification": "
|
10 |
-
}
|
11 |
-
|
12 |
-
# Function to load the selected model
|
13 |
-
@st.cache_resource
|
14 |
-
def load_model(model_name):
|
15 |
-
path = model_paths[model_name]
|
16 |
-
tokenizer = BertTokenizer.from_pretrained(path)
|
17 |
-
model = BertForSequenceClassification.from_pretrained(path)
|
18 |
-
model.eval()
|
19 |
-
return tokenizer, model
|
|
|
1 |
+
from transformers import BertTokenizer, BertForSequenceClassification
|
2 |
+
import streamlit as st
|
3 |
+
|
4 |
+
# Dictionary to map model names to their paths
|
5 |
+
model_paths = {
|
6 |
+
"cahya/bert-base-indonesian-522M": "nlp-brin-id/cahyabert",
|
7 |
+
"indobenchmark/indobert-base-p2": "nlp-brin-id/indobenchmark",
|
8 |
+
"indolem/indobert-base-uncased": "nlp-brin-id/indolem",
|
9 |
+
"mdhugol/indonesia-bert-sentiment-classification": "nlp-brin-id/mdhugol"
|
10 |
+
}
|
11 |
+
|
12 |
+
# Function to load the selected model
|
13 |
+
@st.cache_resource
|
14 |
+
def load_model(model_name):
|
15 |
+
path = model_paths[model_name]
|
16 |
+
tokenizer = BertTokenizer.from_pretrained(path)
|
17 |
+
model = BertForSequenceClassification.from_pretrained(path)
|
18 |
+
model.eval()
|
19 |
+
return tokenizer, model
|