Update app.py

app.py

@@ -1,49 +1,53 @@
 import gradio as gr
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline, AutoConfig
 from huggingface_hub import cached_download, hf_hub_url, list_models
-from transformers.modeling_utils import PreTrainedModel
 import requests
 import json
 import os
 import matplotlib.pyplot as plt
 from io import BytesIO
 import base64
-
-# Choose your backend (PyTorch, TensorFlow, or Flax)
-import torch  # If using PyTorch
+from transformers.models.auto import AutoModel
+from transformers.modeling_utils import PreTrainedModel
+import torch
+from torch.nn.utils import prune
 
 # Function to fetch open-weight LLM models
 def fetch_open_weight_models():
-    models = list_models(filter="open-weight", sort="downloads", limit=10)
-    return [model["id"] for model in models]
+    models = list_models()
+    return models
 
 # Function to prune a model using the "merge-kit" approach
 def prune_model(llm_model_name, target_size, output_dir):
-    # Load the LLM model and tokenizer
-    llm_tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
-    llm_model = AutoModelForSeq2SeqLM.from_pretrained(llm_model_name)
-
-    # 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
-    pruned_model = merge_kit_prune(llm_model, target_num_parameters)
-
-    # Save the pruned model
-    pruned_model.save_pretrained(output_dir)
-
-    # Create a visualization
-    fig, ax = plt.subplots(figsize=(10, 5))
-    ax.bar(["Original", "Pruned"], [config.num_parameters, pruned_model.num_parameters])
-    ax.set_ylabel("Number of Parameters")
-    ax.set_title("Model Size Comparison")
-    buf = BytesIO()
-    fig.savefig(buf, format="png")
-    buf.seek(0)
-    image_base64 = base64.b64encode(buf.read()).decode("utf-8")
-    return f"Pruned model saved to {output_dir}", f"data:image/png;base64,{image_base64}"
+    try:
+        # Load the LLM model and tokenizer
+        llm_tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
+        llm_model = AutoModelForSeq2SeqLM.from_pretrained(llm_model_name)
+
+        # 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
+        pruned_model = merge_kit_prune(llm_model, target_num_parameters)
+
+        # Save the pruned model
+        pruned_model.save_pretrained(output_dir)
+
+        # Create a visualization
+        fig, ax = plt.subplots(figsize=(10, 5))
+        ax.bar(["Original", "Pruned"], [config.num_parameters, pruned_model.num_parameters])
+        ax.set_ylabel("Number of Parameters")
+        ax.set_title("Model Size Comparison")
+        buf = BytesIO()
+        fig.savefig(buf, format="png")
+        buf.seek(0)
+        image_base64 = base64.b64encode(buf.read()).decode("utf-8")
+        return f"Pruned model saved to {output_dir}", f"data:image/png;base64,{image_base64}"
+
+    except Exception as e:
+        return f"Error: {e}", None
 
 # Merge-kit Pruning Function
 def merge_kit_prune(model: PreTrainedModel, target_num_parameters: int) -> PreTrainedModel:
@@ -80,12 +84,8 @@ def create_interface():
     with gr.Blocks() as demo:
         gr.Markdown("## Create a Smaller LLM")
 
-        # Fetch open-weight models from Hugging Face
-        available_models = gr.Dropdown(
-            label="Choose a Large Language Model",
-            choices=fetch_open_weight_models(),
-            interactive=True,
-        )
+        # Input for model name
+        llm_model_name = gr.Textbox(label="Choose a Large Language Model", placeholder="Enter the model name", interactive=True)
 
         # Input for target model size
         target_size = gr.Slider(
@@ -112,7 +112,7 @@ def create_interface():
         # Connect components
         prune_button.click(
             fn=prune_model,
-            inputs=[available_models, target_size, save_model_path],
+            inputs=[llm_model_name, target_size, save_model_path],
             outputs=[pruning_status, visualization],
         )
 
@@ -124,14 +124,17 @@ def create_interface():
         generate_button = gr.Button("Generate Text")
 
         def generate_text(text, model_path):
-            # Load the pruned model and tokenizer
-            tokenizer = AutoTokenizer.from_pretrained(model_path)
-            model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
-
-            # Use the pipeline for text generation
-            generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
-            generated_text = generator(text, max_length=50, num_beams=5, num_return_sequences=1)[0]["generated_text"]
-            return generated_text
+            try:
+                # Load the pruned model and tokenizer
+                tokenizer = AutoTokenizer.from_pretrained(model_path)
+                model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
+
+                # Use the pipeline for text generation
+                generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+                generated_text = generator(text, max_length=50, num_beams=5, num_return_sequences=1)[0]["generated_text"]
+                return generated_text
+            except Exception as e:
+                return f"Error: {e}"
 
         generate_button.click(fn=generate_text, inputs=[text_input, save_model_path], outputs=text_output)