Update app.py

app.py

@@ -47,7 +47,7 @@ def prune_model(llm_model_name, target_size, hf_write_token, repo_name, progress
         target_num_parameters = int(config.num_parameters * (target_size / 100))
 
         # Prune the model
-        pruned_model = merge_kit_prune(llm_model, target_num_parameters, progress)
+        pruned_model = merge_kit_prune(llm_model, target_num_parameters)
 
         log_messages.append("Model pruned successfully.")
         logging.info("Model pruned successfully.")
@@ -81,7 +81,7 @@ def prune_model(llm_model_name, target_size, hf_write_token, repo_name, progress
         return error_message, None, "\n".join(log_messages)
 
 # Merge-kit Pruning Function (adjust as needed)
-def merge_kit_prune(model: PreTrainedModel, target_num_parameters: int, progress) -> PreTrainedModel:
+def merge_kit_prune(model: PreTrainedModel, target_num_parameters: int) -> PreTrainedModel:
     """Prunes a model using a merge-kit approach.
     Args:
         model (PreTrainedModel): The model to be pruned.
@@ -120,9 +120,9 @@ def create_interface():
         pruning_status = gr.Textbox(label="Pruning Status", interactive=False)
         prune_button = gr.Button("Prune Model")
         visualization = gr.Image(label="Model Size Comparison", interactive=False)
-        progress_bar = gr.Progress()
         logs_button = gr.Button("Show Logs")
         logs_output = gr.Textbox(label="Logs", interactive=False)
+        progress_bar = gr.Progress()
 
         def show_logs():
             with open("pruning.log", "r") as log_file:
@@ -131,7 +131,11 @@ def create_interface():
 
         logs_button.click(fn=show_logs, outputs=logs_output)
 
-        prune_button.click(fn=prune_model, inputs=[llm_model_name, target_size, hf_write_token, repo_name, progress_bar], outputs=[pruning_status, visualization, logs_output])
+        def prune_model_with_progress(llm_model_name, target_size, hf_write_token, repo_name):
+            with progress_bar:
+                return prune_model(llm_model_name, target_size, hf_write_token, repo_name)
+
+        prune_button.click(fn=prune_model_with_progress, inputs=[llm_model_name, target_size, hf_write_token, repo_name], outputs=[pruning_status, visualization, logs_output])
 
         text_input = gr.Textbox(label="Input Text")
         text_output = gr.Textbox(label="Generated Text")