Update app.py

app.py

@@ -13,6 +13,7 @@ from torch.nn.utils import prune
 import subprocess
 from tqdm import tqdm
 import logging
+import sys
 
 # Setup logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -47,7 +48,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)
+        pruned_model = merge_kit_prune(llm_model, target_num_parameters, progress)
 
         log_messages.append("Model pruned successfully.")
         logging.info("Model pruned successfully.")
@@ -81,7 +82,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) -> PreTrainedModel:
+def merge_kit_prune(model: PreTrainedModel, target_num_parameters: int, progress: gr.Progress) -> PreTrainedModel:
     """Prunes a model using a merge-kit approach.
     Args:
         model (PreTrainedModel): The model to be pruned.
@@ -100,11 +101,13 @@ def merge_kit_prune(model: PreTrainedModel, target_num_parameters: int) -> PreTr
     for name, module in tqdm(model.named_modules(), desc="Pruning", file=sys.stdout):
         if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)): 
             prune.random_unstructured(module, name="weight", amount=amount)
+            progress(percent_complete=50)  # Example progress update
 
     # Remove the pruned weights
     for name, module in model.named_modules():
         if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)):
             prune.remove(module, name="weight")
+            progress(percent_complete=100)  # Example progress update
 
     return model
 
@@ -132,8 +135,7 @@ def create_interface():
         logs_button.click(fn=show_logs, outputs=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)
+            return prune_model(llm_model_name, target_size, hf_write_token, repo_name, progress_bar)
 
         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])