Update app.py

app.py

@@ -22,14 +22,14 @@ def prune_model(llm_model_name, target_size, output_dir):
     try:
         # Load the LLM model and tokenizer
         llm_tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
-        llm_model = AutoModelForSeq2SeqLM.from_pretrained(llm_model_name)
+        llm_model = AutoModel.from_pretrained(llm_model_name)  # Load using AutoModel
 
         # Get the model config
         config = AutoConfig.from_pretrained(llm_model_name)
         # Calculate the target number of parameters
         target_num_parameters = int(config.num_parameters * (target_size / 100))
 
-        # Use merge-kit to prune the model
+        # Use merge-kit to prune the model (modify pruning logic for Llama)
         pruned_model = merge_kit_prune(llm_model, target_num_parameters)
 
         # Save the pruned model
@@ -67,9 +67,10 @@ def merge_kit_prune(model: PreTrainedModel, target_num_parameters: int) -> PreTr
     # Calculate the pruning amount
     amount = 1 - (target_num_parameters / model.num_parameters)
 
-    # Prune the model using the selected method
+    # Prune the model using the selected method (adapt for Llama)
+    # Example: If Llama uses specific layers, adjust the pruning logic here 
     for name, module in model.named_modules():
-        if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)):
+        if isinstance(module, (torch.nn.Linear, torch.nn.Conv2d)): 
             prune.random_unstructured(module, name="weight", amount=amount)
 
     # Remove the pruned weights