Update app.py

app.py

@@ -12,7 +12,7 @@ if 'model_config' not in st.session_state:
     st.session_state['model_config'] = {}
 
 
-def load_model_config(model_id):
+def load_model_config(model_id, access_token):
     if 'model_id' in st.session_state['model_config'] and st.session_state['model_config']['model_id'] == model_id:
         return st.session_state['model_config']
     if 'parameter_count' in st.session_state:
@@ -41,7 +41,7 @@ def load_model_config(model_id):
         model_config['layernorm_operation'] = 2
 
     try:
-        model_config['model'] = get_model(model_id, None, None)
+        model_config['model'] = get_model(model_id, None, access_token=access_token if len(access_token)>0 else None)
         module_tensors = get_module_tensors(model_config['model'])
         model_config['module_classes'] = classify_module(module_tensors)
     except Exception as e:
@@ -85,7 +85,8 @@ inference_info = {}
 with col1:
     header4("Model")
     model_id = st.text_input("huggingface model id", 'ArthurZ/opt-13b')
-    model_config = load_model_config(model_id)
+    access_token = st.text_input("access token", '')
+    model_config = load_model_config(model_id, access_token)
     model_config['hidden_size'] = st.number_input('hidden size', value=model_config['hidden_size'], format ="%d")
     model_config['num_attention_heads'] = st.number_input('num attention heads', value=model_config['num_attention_heads'], format ="%d")
     model_config['num_hidden_layers'] = st.number_input('num hidden layers', value=model_config['num_hidden_layers'], format ="%d")