update to single model

app/main.py

@@ -5,7 +5,7 @@ try:
     import streamlit as st
     import re
     import streamlit as st
-    from transformers import BertTokenizer, BertModel
+    from transformers import BertTokenizer, AutoConfig
     from model import IndoBERTBiLSTM, IndoBERTModel
 except Exception as e:
     print(e)
@@ -19,19 +19,8 @@ img {
 """
 # Config
 MAX_SEQ_LEN = 128
-# bert_path = 'indolem/indobert-base-uncased'
-bert_path = 'indobenchmark/indobert-base-p2'
-MODELS_PATH = ["kadabengaran/IndoBERT-BiLSTM-Useful-App-Review",
-               "kadabengaran/IndoBERT-Useful-App-Review"]
-
-MODELS_NAME = ["IndoBERT-BiLSTM", "IndoBERT"]
+MODELS_PATH = "kadabengaran/IndoBERT-BiLSTM-Useful-App-Review"
 LABELS = {'Not Useful': 0, 'Useful': 1}
-            
-HIDDEN_DIM = 768
-OUTPUT_DIM = 2 # 2 if Binary
-N_LAYERS = 2
-BIDIRECTIONAL = True
-DROPOUT = 0.2
 
 # Get the Keys
 def get_key(val, my_dict):
@@ -39,19 +28,16 @@ def get_key(val, my_dict):
         if val == value:
             return key
 
-
 def get_device():
     if torch.cuda.is_available():
         return torch.device('cuda')
     else:
         return torch.device('cpu')
 
-
 def load_tokenizer(model_path):
     tokenizer = BertTokenizer.from_pretrained(model_path)
     return tokenizer
 
-
 def remove_special_characters(text):
     text = re.sub(r'[^a-zA-Z0-9\s]', '', text)
     text = re.sub(r"\s+", " ", text)  # replace multiple whitespace characters with a single space
@@ -59,7 +45,6 @@ def remove_special_characters(text):
     text = text.lower()
     return text
 
-
 def preprocess(text, tokenizer, max_seq=MAX_SEQ_LEN):
     return tokenizer.encode_plus(text, add_special_tokens=True, max_length=max_seq,
                                  pad_to_max_length=True,
@@ -68,19 +53,8 @@ def preprocess(text, tokenizer, max_seq=MAX_SEQ_LEN):
                                  )
     
 def load_model():
-    bert = BertModel.from_pretrained(bert_path)
-    
-	# Load the model
-    model_combined = IndoBERTBiLSTM.from_pretrained(MODELS_PATH[0],
-                                     bert,
-                                     HIDDEN_DIM,
-                                     OUTPUT_DIM,
-                                     N_LAYERS, BIDIRECTIONAL,
-                                     DROPOUT)
-    model_base = IndoBERTModel.from_pretrained(MODELS_PATH[1],
-                                     bert,
-                                     OUTPUT_DIM)
-    return model_combined, model_base
+    model = IndoBERTBiLSTM.from_pretrained(MODELS_PATH)
+    return model
 
 def predict_single(text, model, tokenizer, device):
     
@@ -158,22 +132,14 @@ class App:
         st.markdown(html_temp, unsafe_allow_html=True)
         self.render_tabs()
         st.divider()
-        model_choice = self.render_model_selection()
-        if model_choice:
-            if model_choice == MODELS_NAME[0]:
-                model = model_combined
-            elif model_choice == MODELS_NAME[1]:
-                model = model_base
-            self.render_process_button(model, tokenizer, device)
+        model = model_combined
+        self.render_process_button(model, tokenizer, device)
 
     def init_session_state(self):
         if "tab_selected" not in st.session_state:
             st.session_state.tab_selected = tab_labels[0]
 
-    def render_model_selection(self):
-        model_choice = st.selectbox("Select Model", MODELS_NAME)
-        return model_choice
-    
+
     def render_tabs(self):
         tab_selected = st.session_state.get('tab_selected', self.default_tab_selected)
         tab_selected = st.sidebar.radio("Select Input Type", tab_labels)
@@ -195,7 +161,7 @@ class App:
         """
         st.markdown(STYLE, unsafe_allow_html=True)
         file = st.file_uploader("Upload file", type=self.fileTypes)
-
+        # add "untuk kelancaran proses, maksimal csv dengan 1000 baris data"
 
         if not file:
             st.info("Please upload a file of type: " + ", ".join(self.fileTypes))

app/model.py

@@ -1,28 +1,40 @@
 import torch.nn as nn
-from transformers import PreTrainedModel, BertConfig
+import torch
+from transformers import BertModel, BertConfig, PreTrainedModel
+
+def get_device():
+    if torch.cuda.is_available():
+        return torch.device('cuda')
+    else:
+        return torch.device('cpu')
 
 USE_CUDA = False
+device = get_device()
+if device.type == 'cuda':
+    USE_CUDA = True
+
+bert_path = 'indobenchmark/indobert-base-p2'
+HIDDEN_DIM = 768
+OUTPUT_DIM = 2 # 2 if Binary Classification
+N_LAYERS = 1 # 2
+BIDIRECTIONAL = True
+DROPOUT = 0.2 # 0.2
 
 class IndoBERTBiLSTM(PreTrainedModel):
     config_class = BertConfig
-    def __init__(self, bert_config, bert_pretrained_path, hidden_dim, num_classes, n_layers, bidirectional, dropout):
+    def __init__(self, bert_config):
         super().__init__(bert_config)
-        self.output_dim = num_classes
-        self.n_layers = n_layers
-        self.hidden_dim = hidden_dim
-        self.bidirectional = bidirectional
+        self.output_dim = OUTPUT_DIM
+        self.hidden_dim = HIDDEN_DIM
+        self.bidirectional = BIDIRECTIONAL
 
-        self.bert = bert_pretrained_path
+        self.bert = BertModel.from_pretrained(bert_path)
         self.lstm = nn.LSTM(input_size=self.bert.config.hidden_size,
-                            hidden_size=hidden_dim,
-                            num_layers=n_layers,
-                            bidirectional=bidirectional,
+                            hidden_size=self.hidden_dim,
+                            bidirectional=self.bidirectional,
                             batch_first=True)
-        self.dropout = nn.Dropout(dropout)
-        self.global_pooling = nn.AdaptiveAvgPool1d(1)
-        self.hidden_layer = nn.Linear(hidden_dim * 2 if bidirectional else hidden_dim, hidden_dim * 2 if bidirectional else hidden_dim)
-        self.output_layer = nn.Linear(hidden_dim * 2 if bidirectional else hidden_dim, num_classes)
-        self.relu = nn.ReLU()
+        self.dropout = nn.Dropout(DROPOUT)
+        self.output_layer = nn.Linear(self.hidden_dim * 2 if self.bidirectional else self.hidden_dim, self.output_dim)
 
     def forward(self, input_ids, attention_mask):
 
@@ -31,35 +43,26 @@ class IndoBERTBiLSTM(PreTrainedModel):
         output = self.bert(input_ids=input_ids, attention_mask=attention_mask)
         sequence_output = output.last_hidden_state
 
-        # apply dropout
-        sequence_output = self.dropout(sequence_output)
-        # print('output size of the bert:', last_hidden_state.size())
-    
         lstm_output, (hidden_last, cn_last) = self.lstm(sequence_output, hidden)
-        # print('output size of the LSTM:', lstm_output.size())
-        lstm_output = self.dropout(lstm_output)
-
-        # global pooling
-        lstm_output = lstm_output.permute(0, 2, 1)
-        pooled_output = self.global_pooling(lstm_output).squeeze()
+        hidden_last_L=hidden_last[-2]
+        hidden_last_R=hidden_last[-1]
+        hidden_last_out=torch.cat([hidden_last_L,hidden_last_R],dim=-1) #[16, 1536]
 
-        # pass through hidden layer
-        hidden_layer_output = self.hidden_layer(pooled_output)
-        hidden_layer_output = self.relu(hidden_layer_output)
+        # apply dropout
+        out = self.dropout(hidden_last_out)
 
         # output layer
-        logits = self.output_layer(hidden_layer_output)
-        # logits = nn.Softmax(dim=1)(logits)
+        logits = self.output_layer(out)
 
         return logits
 
     def init_hidden(self, batch_size):
         weight = next(self.parameters()).data
-        
+
         number = 1
         if self.bidirectional:
             number = 2
-        
+
         if (USE_CUDA):
             hidden = (weight.new(self.n_layers*number, batch_size, self.hidden_dim).zero_().float().cuda(),
                       weight.new(self.n_layers*number, batch_size, self.hidden_dim).zero_().float().cuda()
@@ -68,21 +71,6 @@ class IndoBERTBiLSTM(PreTrainedModel):
             hidden = (weight.new(self.n_layers*number, batch_size, self.hidden_dim).zero_().float(),
                       weight.new(self.n_layers*number, batch_size, self.hidden_dim).zero_().float()
                      )
-        
-        return hidden
-    
-
-class IndoBERTModel(PreTrainedModel):
-    config_class = BertConfig
-    def __init__(self, bert_config, bert_pretrained, num_classes):
-        super().__init__(bert_config)
-        self.bert = bert_pretrained
-        self.dropout = nn.Dropout(0.1)
-        self.fc = nn.Linear(self.bert.config.hidden_size, num_classes)
 
-    def forward(self, input_ids, attention_mask):
-        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
-        pooled_output = outputs.pooler_output
-        pooled_output = self.dropout(pooled_output)
-        logits = self.fc(pooled_output)
-        return logits
+        return hidden
+