debug

app.py

@@ -49,9 +49,9 @@ def load_css(file_name):
 
 @st.cache(show_spinner=True,allow_output_mutation=True)
 def load_model():
-    tokenizer = AlbertTokenizer.from_pretrained('albert-xxlarge-v2')
+    tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
     #model = CustomFlaxAlbertForMaskedLM.from_pretrained('albert-xxlarge-v2',from_pt=True)
-    model = AlbertForMaskedLM.from_pretrained('albert-xxlarge-v2')
+    model = AlbertForMaskedLM.from_pretrained('albert-base-v2')
     return tokenizer,model
 
 def clear_data():
@@ -163,7 +163,8 @@ if __name__=='__main__':
     wide_setup()
     #load_css('style.css')
     tokenizer,model = load_model()
-    num_layers, num_heads = 12, 64
+    num_layers, num_heads = model.config.num_hidden_layers, model.config.num_attention_heads
+    st.write(num_layers,num_heads)
     mask_id = tokenizer('[MASK]').input_ids[1:-1][0]
 
     main_area = st.empty()