Update app.py

app.py

@@ -325,6 +325,10 @@ with gr.Blocks() as demo:
             """)
             with gr.Row():
                 with gr.Column():
+                    hf_model_name_or_path = gr.Textbox(
+                        label="HuggingFace Model Name or Path",
+                        info="Name of the HuggingFace Hub repository or the local file path for it"
+                    )
                     vocab_size = gr.Number(
                         label="Vocab Size", 
                         value=51200, 
@@ -350,25 +354,10 @@ with gr.Blocks() as demo:
                         value=1.0, 
                         info="Ratio of kv heads to query heads used in model. 1.0 for MHA"
                     )
-                    topk = gr.Number(
-                        label="Top K Routing for MoE", 
-                        value=1, 
-                        info="Top k routing for Mixture of Experts (MoE)"
-                    )
-                    moe = gr.Checkbox(
-                        label="Mixture of Experts (MoE)", 
-                        value=False, 
-                        info="Whether the model uses Mixture of Experts"
-                    )
-                    num_experts = gr.Number(
-                        label="Number of Experts", 
-                        value=128, 
-                        info="Number of experts for Mixture of Experts (MoE)"
-                    )
-                    expert_interval = gr.Number(
-                        label="Expert Interval", 
-                        value=2, 
-                        info="Expert interval for Mixture of Experts (MoE)"
+                    ffn_expansion_factor = gr.Number(
+                        label="FFN Expansion Factor", 
+                        value=4, 
+                        info="How much the MLP hidden size expands"
                     )
                     batch_size = gr.Number(
                         label="Batch Size", 
@@ -385,17 +374,35 @@ with gr.Blocks() as demo:
                         value=True, 
                         info="Whether Megatron-style activation checkpointing is being used"
                     )
-                    ffn_expansion_factor = gr.Number(
-                        label="FFN Expansion Factor", 
-                        value=4, 
-                        info="How much the MLP hidden size expands"
-                    )
                     infer = gr.Checkbox(
                         label="Inference-Only", 
                         value=False, 
                         info="Whether the model is being used for inference-only"
                     )
 
+                    # MoE parameters hidden in accordion
+                    with gr.Accordion("Mixture of Experts (MoE)", open=False):
+                        moe = gr.Checkbox(
+                            label="Mixture of Experts (MoE)", 
+                            value=False, 
+                            info="Whether the model uses Mixture of Experts"
+                        )
+                        num_experts = gr.Number(
+                            label="Number of Experts", 
+                            value=128, 
+                            info="Number of experts for Mixture of Experts (MoE)"
+                        )
+                        expert_interval = gr.Number(
+                            label="Expert Interval", 
+                            value=2, 
+                            info="Expert interval for Mixture of Experts (MoE)"
+                        )
+                        topk = gr.Number(
+                            label="Top K Routing for MoE", 
+                            value=1, 
+                            info="Top k routing for Mixture of Experts (MoE)"
+                        )
+
             calc_flops_button = gr.Button("Calculate FLOPs")
             flops_result = gr.JSON(label="FLOP Calculation Result", interactive=False)
             calc_flops_button.click(
@@ -404,8 +411,8 @@ with gr.Blocks() as demo:
                 outputs=flops_result
             )
 
-            hf_model_name_or_path = gr.Textbox(label="HuggingFace Model Name or Path", info="Name of the HuggingFace model or local path")
-            hf_model_name_or_path.change(fn=get_hf_model_args, inputs=[hf_model_name_or_path], outputs=[num_layers, hidden_size, vocab_size, sequence_length])
-
+            hf_model_name_or_path.change(fn=update_from_hf_model, 
+                inputs=[hf_model_name_or_path], 
+                outputs=[num_layers, hidden_size, vocab_size, sequence_length])
 
 demo.launch()