Update app.py

app.py

@@ -10,32 +10,37 @@ from io import BytesIO
 import base64
 import torch
 from torch.nn.utils import prune
+import subprocess
 
 # Function to fetch open-weight LLM models
 def fetch_open_weight_models():
     models = list_models()
     return models
 
+# Ensure sentencepiece is installed
+try:
+    import sentencepiece
+except ImportError:
+    subprocess.check_call(["pip", "install", "sentencepiece"])
+
 # Function to prune a model using the "merge-kit" approach
 def prune_model(llm_model_name, target_size, hf_write_token, repo_name):
     try:
         # Load the LLM model and tokenizer
         llm_tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
-        # Handle cases where the model is split into multiple safetensors
         llm_model = AutoModelForCausalLM.from_pretrained(
             llm_model_name,
-            torch_dtype=torch.float16,  # Adjust dtype as needed
+            torch_dtype=torch.float16,
         )
 
         # 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
+        # Prune the model
         pruned_model = merge_kit_prune(llm_model, target_num_parameters)
 
-        # Save the pruned model to Hugging Face repository
+        # Save the pruned model
         api = HfApi()
         repo_id = f"{hf_write_token}/{repo_name}"
         create_repo(repo_id, token=hf_write_token, private=False, exist_ok=True)
@@ -51,10 +56,11 @@ def prune_model(llm_model_name, target_size, hf_write_token, repo_name):
         fig.savefig(buf, format="png")
         buf.seek(0)
         image_base64 = base64.b64encode(buf.read()).decode("utf-8")
-        return f"Pruned model saved to Hugging Face Hub in repository {repo_id}", f"data:image/png;base64,{image_base64}"
+
+        return f"Pruned model saved to Hugging Face Hub in repository {repo_id}", f"data:image/png;base64,{image_base64}", None
 
     except Exception as e:
-        return f"Error: {e}", None
+        return f"Error: {e}", None, None
 
 # Merge-kit Pruning Function (adjust as needed)
 def merge_kit_prune(model: PreTrainedModel, target_num_parameters: int) -> PreTrainedModel:
@@ -88,55 +94,24 @@ def create_interface():
     with gr.Blocks() as demo:
         gr.Markdown("## Create a Smaller LLM")
 
-        # 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(
-            label="Target Model Size (%)",
-            minimum=1,
-            maximum=100,
-            step=1,
-            value=50,
-            interactive=True,
-        )
-
-        # Input for Hugging Face write token
+        target_size = gr.Slider(label="Target Model Size (%)", minimum=1, maximum=100, step=1, value=50, interactive=True)
         hf_write_token = gr.Textbox(label="Hugging Face Write Token", placeholder="Enter your HF write token", interactive=True, type="password")
-
-        # Input for repository name
         repo_name = gr.Textbox(label="Repository Name", placeholder="Enter the name of the repository", interactive=True)
-
-        # Output for pruning status
         pruning_status = gr.Textbox(label="Pruning Status", interactive=False)
-
-        # Button to start pruning
         prune_button = gr.Button("Prune Model")
-
-        # Output for visualization
         visualization = gr.Image(label="Model Size Comparison", interactive=False)
 
-        # Connect components
-        prune_button.click(
-            fn=prune_model,
-            inputs=[llm_model_name, target_size, hf_write_token, repo_name],
-            outputs=[pruning_status, visualization],
-        )
+        prune_button.click(fn=prune_model, inputs=[llm_model_name, target_size, hf_write_token, repo_name], outputs=[pruning_status, visualization])
 
-        # Example usage of the pruned model (optional)
         text_input = gr.Textbox(label="Input Text")
         text_output = gr.Textbox(label="Generated Text")
-
-        # Generate text button
         generate_button = gr.Button("Generate Text")
 
         def generate_text(text, repo_name):
             try:
-                # Load the pruned model and tokenizer
                 tokenizer = AutoTokenizer.from_pretrained(repo_name, use_auth_token=hf_write_token)
                 model = AutoModelForCausalLM.from_pretrained(repo_name, use_auth_token=hf_write_token)
-
-                # 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
@@ -149,4 +124,4 @@ def create_interface():
 
 # Create and launch the Gradio interface
 demo = create_interface()
-demo.launch(share=True)
+demo.launch(share=True)