Update app.py

app.py

@@ -1,51 +1,251 @@
 import gradio as gr
-from transformers import AutoModel
+from transformers import AutoModel, AutoConfig
 import torch
+import json
+from collections import defaultdict, OrderedDict
 
-def count_parameters(model_path):
+def analyze_model_parameters(model_path, show_layer_details=False):
     try:
+        # Load model configuration first
+        config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+        
         # Load model on CPU
         model = AutoModel.from_pretrained(model_path, device_map="cpu", trust_remote_code=True)
         
-        # Count trainable parameters (accounting for weight tying)
+        # Initialize counters
+        total_params = 0
+        trainable_params = 0
+        embedding_params = 0
+        non_embedding_params = 0
+        
+        # Track unique parameters to handle weight tying
         unique_params = {}
-        for name, p in model.named_parameters():
-            if p.requires_grad:
-                unique_params[p.data_ptr()] = (name, p.numel())
+        param_details = []
+        layer_breakdown = defaultdict(lambda: {'total': 0, 'trainable': 0, 'params': []})
         
-        trainable_params = sum(numel for _, numel in unique_params.values())
+        # Embedding layer patterns (common names for embedding layers)
+        embedding_patterns = [
+            'embeddings', 'embed', 'wte', 'wpe', 'word_embedding', 
+            'position_embedding', 'token_embedding', 'embed_tokens',
+            'embed_positions', 'embed_layer_norm'
+        ]
         
-        # Count total parameters (accounting for weight tying)
-        unique_params = {}
-        for name, p in model.named_parameters():
-            unique_params[p.data_ptr()] = (name, p.numel())
+        def is_embedding_param(name):
+            name_lower = name.lower()
+            return any(pattern in name_lower for pattern in embedding_patterns)
+        
+        def get_layer_name(param_name):
+            """Extract layer information from parameter name"""
+            parts = param_name.split('.')
+            if len(parts) >= 2:
+                # Handle common transformer architectures
+                if 'layer' in parts or 'layers' in parts:
+                    for i, part in enumerate(parts):
+                        if part in ['layer', 'layers'] and i + 1 < len(parts):
+                            try:
+                                layer_num = int(parts[i + 1])
+                                return f"Layer {layer_num}"
+                            except ValueError:
+                                pass
+                # Handle other patterns
+                if 'encoder' in parts:
+                    return "Encoder"
+                elif 'decoder' in parts:
+                    return "Decoder"
+                elif any(emb in param_name.lower() for emb in embedding_patterns):
+                    return "Embeddings"
+                elif 'classifier' in param_name.lower() or 'head' in param_name.lower():
+                    return "Classification Head"
+                elif 'pooler' in param_name.lower():
+                    return "Pooler"
+                elif 'ln' in param_name.lower() or 'norm' in param_name.lower():
+                    return "Layer Norm"
+            return "Other"
+        
+        # Analyze all parameters
+        for name, param in model.named_parameters():
+            param_size = param.numel()
+            is_trainable = param.requires_grad
+            is_embedding = is_embedding_param(name)
+            layer_name = get_layer_name(name)
+            
+            # Handle weight tying by using data pointer
+            ptr = param.data_ptr()
+            if ptr not in unique_params:
+                unique_params[ptr] = {
+                    'name': name,
+                    'size': param_size,
+                    'trainable': is_trainable,
+                    'embedding': is_embedding,
+                    'layer': layer_name,
+                    'shape': list(param.shape)
+                }
+                
+                # Add to totals
+                total_params += param_size
+                if is_trainable:
+                    trainable_params += param_size
+                if is_embedding:
+                    embedding_params += param_size
+                else:
+                    non_embedding_params += param_size
+                    
+                # Add to layer breakdown
+                layer_breakdown[layer_name]['total'] += param_size
+                if is_trainable:
+                    layer_breakdown[layer_name]['trainable'] += param_size
+                
+            # Add parameter details
+            param_details.append({
+                'name': name,
+                'shape': list(param.shape),
+                'size': param_size,
+                'trainable': is_trainable,
+                'embedding': is_embedding,
+                'layer': layer_name,
+                'shared': ptr in [p['ptr'] for p in param_details if 'ptr' in p],
+                'ptr': ptr
+            })
+            
+            # Add to layer breakdown details
+            layer_breakdown[layer_name]['params'].append({
+                'name': name,
+                'shape': list(param.shape),
+                'size': param_size,
+                'trainable': is_trainable
+            })
+        
+        # Format the summary
+        summary = f"""
+🔍 **MODEL ANALYSIS: {model_path}**
+
+📊 **PARAMETER SUMMARY**
+├── Total Parameters: {total_params:,}
+├── Trainable Parameters: {trainable_params:,}
+├── Non-trainable Parameters: {total_params - trainable_params:,}
+└── Trainable Percentage: {(trainable_params/total_params*100):.1f}%
+
+🧠 **PARAMETER BREAKDOWN**
+├── Embedding Parameters: {embedding_params:,} ({embedding_params/total_params*100:.1f}%)
+└── Non-embedding Parameters: {non_embedding_params:,} ({non_embedding_params/total_params*100:.1f}%)
+
+📋 **MODEL INFO**
+├── Model Type: {config.model_type if hasattr(config, 'model_type') else 'Unknown'}
+├── Architecture: {config.architectures[0] if hasattr(config, 'architectures') and config.architectures else 'Unknown'}
+└── Hidden Size: {getattr(config, 'hidden_size', 'Unknown')}
+"""
+
+        # Add layer breakdown summary
+        if layer_breakdown:
+            summary += "\n🏗️ **LAYER BREAKDOWN SUMMARY**\n"
+            sorted_layers = sorted(layer_breakdown.items(), key=lambda x: (
+                0 if x[0] == "Embeddings" else
+                1 if x[0].startswith("Layer") else
+                2 if x[0] == "Layer Norm" else
+                3 if x[0] == "Pooler" else
+                4 if x[0] == "Classification Head" else 5
+            ))
+            
+            for layer_name, info in sorted_layers:
+                percentage = info['total'] / total_params * 100
+                summary += f"├── {layer_name}: {info['total']:,} params ({percentage:.1f}%)\n"
+        
+        # Detailed layer breakdown if requested
+        layer_details = ""
+        if show_layer_details:
+            layer_details = "\n" + "="*60 + "\n"
+            layer_details += "🔍 **DETAILED LAYER-BY-LAYER BREAKDOWN**\n"
+            layer_details += "="*60 + "\n"
+            
+            for layer_name, info in sorted_layers:
+                layer_details += f"\n📁 **{layer_name.upper()}**\n"
+                layer_details += f"   Total: {info['total']:,} | Trainable: {info['trainable']:,}\n"
+                layer_details += f"   Parameters:\n"
+                
+                for param_info in info['params']:
+                    trainable_mark = "✓" if param_info['trainable'] else "✗"
+                    layer_details += f"   {trainable_mark} {param_info['name']}: {param_info['shape']} → {param_info['size']:,}\n"
         
-        total_params = sum(numel for _, numel in unique_params.values())
+        return summary + layer_details
         
-        # Format numbers with commas for readability
-        return f"""
-Total Parameters: {total_params:,}
-Trainable Parameters: {trainable_params:,}
-        """
     except Exception as e:
-        return f"Error loading model: {str(e)}"
-
-# Create Gradio interface
-demo = gr.Interface(
-    fn=count_parameters,
-    inputs=gr.Textbox(
-        label="Enter Hugging Face Model Path",
-        placeholder="e.g., bert-base-uncased"
-    ),
-    outputs=gr.Textbox(label="Parameter Count"),
-    title="Hugging Face Model Parameter Counter",
-    description="Enter a Hugging Face model path to see its parameter count.",
-    examples=[
-        ["bert-base-uncased"],
-        ["gpt2"],
-        ["roberta-base"]
-    ]
-)
+        return f"❌ **Error loading model:** {str(e)}\n\nPlease check that the model path is correct and the model is accessible."
+
+def count_parameters_basic(model_path):
+    """Basic parameter counting without layer details"""
+    return analyze_model_parameters(model_path, show_layer_details=False)
+
+def count_parameters_detailed(model_path):
+    """Detailed parameter counting with layer-by-layer breakdown"""
+    return analyze_model_parameters(model_path, 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:
+    gr.Markdown("""
+    # 🤗 Advanced HuggingFace Model Parameter Analyzer
+    
+    Enter any HuggingFace model path to get detailed parameter analysis including:
+    - **Total & trainable parameter counts**
+    - **Embedding vs non-embedding breakdown** 
+    - **Layer-by-layer analysis**
+    - **Weight sharing detection**
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            model_input = gr.Textbox(
+                label="🔍 HuggingFace Model Path",
+                placeholder="e.g., bert-base-uncased, gpt2, microsoft/DialoGPT-medium",
+                value="bert-base-uncased"
+            )
+            
+        with gr.Column(scale=1):
+            analyze_btn = gr.Button("📊 Analyze Model", variant="primary")
+            detailed_btn = gr.Button("🔍 Detailed Analysis", variant="secondary")
+    
+    output_text = gr.Textbox(
+        label="📋 Analysis Results",
+        lines=20,
+        max_lines=50,
+        show_copy_button=True
+    )
+    
+    # Event handlers
+    analyze_btn.click(
+        fn=count_parameters_basic,
+        inputs=model_input,
+        outputs=output_text
+    )
+    
+    detailed_btn.click(
+        fn=count_parameters_detailed,
+        inputs=model_input,
+        outputs=output_text
+    )
+    
+    # Example models
+    gr.Examples(
+        examples=[
+            ["bert-base-uncased"],
+            ["gpt2"],
+            ["roberta-base"],
+            ["distilbert-base-uncased"],
+            ["microsoft/DialoGPT-medium"],
+            ["facebook/bart-base"],
+            ["t5-small"],
+            ["google/flan-t5-small"]
+        ],
+        inputs=model_input,
+        label="🎯 Example Models"
+    )
+    
+    gr.Markdown("""
+    ### 📝 Notes:
+    - **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
+    - **Model compatibility**: Works with most HuggingFace transformer models
+    """)
 
 if __name__ == "__main__":
     demo.launch()