Update app.py

app.py

@@ -8,7 +8,6 @@ import base64
 import torch
 from torch.nn.utils import prune
 import subprocess
-from tqdm import tqdm
 import logging
 import sys
 
@@ -46,7 +45,6 @@ def merge_kit_prune(model: PreTrainedModel, target_num_parameters: int, progress
     Returns:
         PreTrainedModel: The pruned model.
     """
-    
     total_params = sum(p.numel() for p in model.parameters())
     amount = 1 - (target_num_parameters / total_params)
 
@@ -93,13 +91,12 @@ def prune_model(llm_model_name, target_size, hf_write_token, repo_name, base_mod
         
         # 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)
-        pruned_model.push_to_hub(repo_id, use_auth_token=hf_write_token)
-        llm_tokenizer.push_to_hub(repo_id, use_auth_token=hf_write_token)
+        create_repo(repo_name, token=hf_write_token, private=False, exist_ok=True)
+        pruned_model.push_to_hub(repo_name, use_auth_token=hf_write_token)
+        llm_tokenizer.push_to_hub(repo_name, use_auth_token=hf_write_token)
         
-        log_messages.append(f"Pruned model saved to Hugging Face Hub in repository {repo_id}")
-        logging.info(f"Pruned model saved to Hugging Face Hub in repository {repo_id}")
+        log_messages.append(f"Pruned model saved to Hugging Face Hub in repository {repo_name}")
+        logging.info(f"Pruned model saved to Hugging Face Hub in repository {repo_name}")
         
         # Create a visualization
         fig, ax = plt.subplots(figsize=(10, 5))
@@ -111,7 +108,7 @@ def prune_model(llm_model_name, target_size, hf_write_token, repo_name, base_mod
         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}", '\n'.join(log_messages)
+        return f"Pruned model saved to Hugging Face Hub in repository {repo_name}", f"data:image/png;base64,{image_base64}", '\n'.join(log_messages)
     
     except Exception as e:
         error_message = f"Detailed error: {repr(e)}"
@@ -119,6 +116,18 @@ def prune_model(llm_model_name, target_size, hf_write_token, repo_name, base_mod
         logging.error(error_message)
         return error_message, None, '\n'.join(log_messages)
 
+# Define function to generate text
+def generate_text(text, repo_name, hf_write_token):
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(repo_name, use_auth_token=hf_write_token)
+        model = AutoModelForCausalLM.from_pretrained(repo_name, use_auth_token=hf_write_token)
+        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:
+        logging.error(f"Error during text generation: {e}")
+        return f"Error: {repr(e)}"
+
 # Function to create a Gradio interface
 def create_interface():
     with gr.Blocks() as demo:
@@ -126,43 +135,33 @@ def create_interface():
         
         # Fetch available model names
         model_names = get_model_names()
-        
+
         # Input components
         llm_model_name = gr.Dropdown(label="Choose a Large Language Model", choices=model_names, interactive=True)
+        base_model_name = gr.Dropdown(label="Base Model Name (if required)", choices=model_names, interactive=True, visible=False)
         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")
         repo_name = gr.Textbox(label="Repository Name", placeholder="Enter the name of the repository", interactive=True)
         pruned_func_choice = gr.Radio(label="Pruning Function", choices=["merge-kit"], value="merge-kit", interactive=True)
-        base_model_name = gr.Dropdown(label="Base Model Name (if required)", choices=model_names, interactive=True, visible=False)
         
         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()
-        
-        def prune_model_with_progress(llm_model_name, target_size, hf_write_token, repo_name, pruned_func_choice, base_model_name):
+
+        # Define function for pruning model with progress
+        def prune_model_with_progress(llm_model_name, base_model_name, target_size, hf_write_token, repo_name, pruned_func_choice):
             if pruned_func_choice == "merge-kit":
                 return prune_model(llm_model_name, target_size, hf_write_token, repo_name, base_model_name, progress_bar)
             else:
                 return f"Pruning function '{pruned_func_choice}' not implemented.", None, None
-        
-        prune_button.click(fn=prune_model_with_progress, inputs=[llm_model_name, target_size, hf_write_token, repo_name, pruned_func_choice, base_model_name], outputs=[pruning_status, visualization])
-        
+
+        prune_button.click(fn=prune_model_with_progress, inputs=[llm_model_name, base_model_name, target_size, hf_write_token, repo_name, pruned_func_choice], outputs=[pruning_status, visualization])
+
         text_input = gr.Textbox(label="Input Text")
         text_output = gr.Textbox(label="Generated Text")
         generate_button = gr.Button("Generate Text")
-        
-        def generate_text(text, repo_name, hf_write_token):
-            try:
-                tokenizer = AutoTokenizer.from_pretrained(repo_name, use_auth_token=hf_write_token)
-                model = AutoModelForCausalLM.from_pretrained(repo_name, use_auth_token=hf_write_token)
-                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:
-                logging.error(f"Error during text generation: {e}")
-                return f"Error: {repr(e)}"
-        
+
         generate_button.click(fn=generate_text, inputs=[text_input, repo_name, hf_write_token], outputs=text_output)
 
     return demo