Update app.py

app.py

@@ -4,13 +4,18 @@ import torch
 import json
 from collections import defaultdict, OrderedDict
 
-def analyze_model_parameters(model_path, show_layer_details=False):
+def analyze_model_parameters(model_path, hf_token=None, show_layer_details=False):
     try:
+        # Prepare token parameter
+        token_kwargs = {}
+        if hf_token and hf_token.strip():
+            token_kwargs['token'] = hf_token.strip()
+        
         # Load model configuration first
-        config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+        config = AutoConfig.from_pretrained(model_path, trust_remote_code=True, **token_kwargs)
         
         # Load model on CPU
-        model = AutoModel.from_pretrained(model_path, device_map="cpu", trust_remote_code=True)
+        model = AutoModel.from_pretrained(model_path, device_map="cpu", trust_remote_code=True, **token_kwargs)
         
         # Initialize counters
         total_params = 0
@@ -169,15 +174,21 @@ def analyze_model_parameters(model_path, show_layer_details=False):
         return summary + layer_details
         
     except Exception as e:
-        return f"❌ **Error loading model:** {str(e)}\n\nPlease check that the model path is correct and the model is accessible."
+        error_msg = str(e)
+        if "401" in error_msg or "authentication" in error_msg.lower():
+            return f"🔒 **Authentication Error:** This model requires a valid HuggingFace token.\n\nPlease provide your HuggingFace token in the token field above.\n\nOriginal error: {error_msg}"
+        elif "404" in error_msg or "not found" in error_msg.lower():
+            return f"🔍 **Model Not Found:** The model '{model_path}' was not found.\n\nPlease check:\n- Model path is correct\n- Model exists on HuggingFace Hub\n- You have access to the model (use token if private)\n\nOriginal error: {error_msg}"
+        else:
+            return f"❌ **Error loading model:** {error_msg}\n\nPlease check that the model path is correct and accessible."
 
-def count_parameters_basic(model_path):
+def count_parameters_basic(model_path, hf_token=None):
     """Basic parameter counting without layer details"""
-    return analyze_model_parameters(model_path, show_layer_details=False)
+    return analyze_model_parameters(model_path, hf_token, show_layer_details=False)
 
-def count_parameters_detailed(model_path):
+def count_parameters_detailed(model_path, hf_token=None):
     """Detailed parameter counting with layer-by-layer breakdown"""
-    return analyze_model_parameters(model_path, show_layer_details=True)
+    return analyze_model_parameters(model_path, hf_token, show_layer_details=True)
 
 # Create Gradio interface with multiple outputs
 with gr.Blocks(title="🤗 Advanced HuggingFace Model Parameter Analyzer", theme=gr.themes.Soft()) as demo:
@@ -189,6 +200,7 @@ with gr.Blocks(title="🤗 Advanced HuggingFace Model Parameter Analyzer", theme
     - **Embedding vs non-embedding breakdown** 
     - **Layer-by-layer analysis**
     - **Weight sharing detection**
+    - **Private model access** with HuggingFace token
     """)
     
     with gr.Row():
@@ -200,8 +212,16 @@ with gr.Blocks(title="🤗 Advanced HuggingFace Model Parameter Analyzer", theme
             )
             
         with gr.Column(scale=1):
-            analyze_btn = gr.Button("📊 Analyze Model", variant="primary")
-            detailed_btn = gr.Button("🔍 Detailed Analysis", variant="secondary")
+            hf_token_input = gr.Textbox(
+                label="🔑 HuggingFace Token (Optional)",
+                placeholder="hf_...",
+                type="password",
+                info="Required for private models or gated models"
+            )
+            
+    with gr.Row():
+        analyze_btn = gr.Button("📊 Analyze Model", variant="primary")
+        detailed_btn = gr.Button("🔍 Detailed Analysis", variant="secondary")
     
     output_text = gr.Textbox(
         label="📋 Analysis Results",
@@ -213,13 +233,13 @@ with gr.Blocks(title="🤗 Advanced HuggingFace Model Parameter Analyzer", theme
     # Event handlers
     analyze_btn.click(
         fn=count_parameters_basic,
-        inputs=model_input,
+        inputs=[model_input, hf_token_input],
         outputs=output_text
     )
     
     detailed_btn.click(
         fn=count_parameters_detailed,
-        inputs=model_input,
+        inputs=[model_input, hf_token_input],
         outputs=output_text
     )
     
@@ -244,6 +264,8 @@ with gr.Blocks(title="🤗 Advanced HuggingFace Model Parameter Analyzer", theme
     - **Weight tying detection**: Automatically handles shared parameters (e.g., input/output embeddings)
     - **Layer categorization**: Groups parameters by transformer layers, embeddings, etc.
     - **Detailed analysis**: Click "Detailed Analysis" for parameter-by-parameter breakdown
+    - **Private models**: Use your HuggingFace token to access private or gated models
+    - **Token security**: Token is only used for this session and not stored
     - **Model compatibility**: Works with most HuggingFace transformer models
     """)