model type

app.py

@@ -47,10 +47,10 @@ def load_css(file_name):
         st.markdown(f'<style>{f.read()}</style>', unsafe_allow_html=True)
 
 @st.cache(show_spinner=True,allow_output_mutation=True)
-def load_model():
-    tokenizer = AlbertTokenizer.from_pretrained('albert-xxlarge-v2')
+def load_model(model_name):
+    tokenizer = AlbertTokenizer.from_pretrained(model_name)
     #model = CustomFlaxAlbertForMaskedLM.from_pretrained('albert-xxlarge-v2',from_pt=True)
-    model = AlbertForMaskedLM.from_pretrained('albert-xxlarge-v2')
+    model = AlbertForMaskedLM.from_pretrained(model_name)
     return tokenizer,model
 
 def clear_data():
@@ -167,14 +167,23 @@ def run_intervention(interventions,batch_size,model,masked_ids_option_1,masked_i
 if __name__=='__main__':
     wide_setup()
     load_css('style.css')
-    tokenizer,model = load_model()
-    num_layers, num_heads = model.config.num_hidden_layers, model.config.num_attention_heads
-    mask_id = tokenizer('[MASK]').input_ids[1:-1][0]
-
-    main_area = st.empty()
 
     if 'page_status' not in st.session_state:
-        st.session_state['page_status'] = 'type_in'
+        st.session_state['page_status'] = 'model_selection'
+
+    if st.session_state['page_status']=='model_selection':
+        model_name = st.selectbox('Please select the model from below.',
+                                    ('bert-base-uncased','bert-large-cased',
+                                    'roberta-base','roberta-large',
+                                    'albert-base-v2','albert-large-v2','albert-xlarge-v2','albert-xxlarge-v2'),index=3)
+        st.sesstion_state['model_name'] = model_name
+        if st.button('Confirm',key='model_name'):
+            st.session_state['page_status'] = 'type_in'
+            st.experimental_rerun()
+
+    tokenizer,model = load_model(st.session_state['model_name'])
+    num_layers, num_heads = model.config.num_hidden_layers, model.config.num_attention_heads
+    mask_id = tokenizer('[MASK]').input_ids[1:-1][0]
 
     if st.session_state['page_status']=='type_in':
         show_instruction('1. Type in the sentences and click "Tokenize"',fontsize=16)