Update app.py

app.py

@@ -22,7 +22,7 @@ logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(
 try:
     import sentencepiece
 except ImportError:
-    subprocess.check_call(["pip", "install", "sentencepiece"])
+    subprocess.check_call(['pip', 'install', 'sentencepiece'])
 
 # Function to fetch open-weight LLM models
 def fetch_open_weight_models():
@@ -40,8 +40,8 @@ def prune_model(llm_model_name, target_size, hf_write_token, repo_name, progress
             torch_dtype=torch.float16,
         )
 
-        log_messages.append("Model and tokenizer loaded successfully.")
-        logging.info("Model and tokenizer loaded successfully.")
+        log_messages.append('Model and tokenizer loaded successfully.')
+        logging.info('Model and tokenizer loaded successfully.')
 
         # Get the model config
         config = AutoConfig.from_pretrained(llm_model_name)
@@ -50,8 +50,8 @@ def prune_model(llm_model_name, target_size, hf_write_token, repo_name, progress
         # Prune the model
         pruned_model = merge_kit_prune(llm_model, target_num_parameters, progress)
 
-        log_messages.append("Model pruned successfully.")
-        logging.info("Model pruned successfully.")
+        log_messages.append('Model pruned successfully.')
+        logging.info('Model pruned successfully.')
 
         # Save the pruned model
         api = HfApi()
@@ -65,21 +65,21 @@ def prune_model(llm_model_name, target_size, hf_write_token, repo_name, progress
 
         # Create a visualization
         fig, ax = plt.subplots(figsize=(10, 5))
-        ax.bar(["Original", "Pruned"], [config.num_parameters, sum(p.numel() for p in pruned_model.parameters())])
-        ax.set_ylabel("Number of Parameters")
-        ax.set_title("Model Size Comparison")
+        ax.bar(['Original', 'Pruned'], [config.num_parameters, sum(p.numel() for p in pruned_model.parameters())])
+        ax.set_ylabel('Number of Parameters')
+        ax.set_title('Model Size Comparison')
         buf = BytesIO()
-        fig.savefig(buf, format="png")
+        fig.savefig(buf, format='png')
         buf.seek(0)
-        image_base64 = base64.b64encode(buf.read()).decode("utf-8")
+        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_id}", f"data:image/png;base64,{image_base64}", '\n'.join(log_messages)
 
     except Exception as e:
         error_message = f"Error: {e}"
         log_messages.append(error_message)
         logging.error(error_message)
-        return error_message, None, "\n".join(log_messages)
+        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, progress: gr.Progress) -> PreTrainedModel:
@@ -90,25 +90,20 @@ def merge_kit_prune(model: PreTrainedModel, target_num_parameters: int, progress
     Returns:
         PreTrainedModel: The pruned model.
     """
-    # Define the pruning method
-    pruning_method = "unstructured"
-
-    # Calculate the pruning amount
+    
     total_params = sum(p.numel() for p in model.parameters())
-    amount = 1 - (target_num_parameters / total_params)
+    amount = 1 - (target_num_parameters / total_params) 
 
-    # Prune the model using the selected method
-    for name, module in tqdm(model.named_modules(), desc="Pruning", file=sys.stdout):
+    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
+            prune.random_unstructured(module, name='weight', amount=amount)
+            progress(percent_complete=50) 
 
-    # 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
-
+            prune.remove(module, name='weight')
+            progress(percent_complete=100)
+    
     return model
 
 # Function to create a Gradio interface
@@ -128,7 +123,7 @@ def create_interface():
         progress_bar = gr.Progress()
 
         def show_logs():
-            with open("pruning.log", "r") as log_file:
+            with open('pruning.log', 'r') as log_file:
                 logs = log_file.read()
             return logs
 
@@ -143,20 +138,20 @@ def create_interface():
         text_output = gr.Textbox(label="Generated Text")
         generate_button = gr.Button("Generate Text")
 
-        def generate_text(text, repo_name):
+        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"]
+                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, repo_name], outputs=text_output)
+        generate_button.click(fn=generate_text, inputs=[text_input, repo_name, hf_write_token], outputs=text_output)
 
     return demo
 
 # Create and launch the Gradio interface
 demo = create_interface()
-demo.launch(share=True)
+demo.launch()