attribute error

scripts/evalexperts.py

@@ -77,7 +77,7 @@ class ExpertTrackingHFLM(HFLM):
             )
             
             # Update statistics
-            self._update_expert_stats(
+            self.update_expert_stats(  # Changed from _update_expert_stats to update_expert_stats
                 layer_idx=layer_idx,
                 topk_experts=topk_experts,
                 topk_probs=topk_probs,
@@ -89,6 +89,61 @@ class ExpertTrackingHFLM(HFLM):
             
         return expert_hook
     
+    def update_expert_stats(self, layer_idx: int, topk_experts: torch.Tensor,  # Renamed from _update_expert_stats
+                          topk_probs: torch.Tensor, num_regular_experts: int,
+                          num_small_experts: int, batch_size: int, seq_len: int):
+        """Update expert usage statistics."""
+        # Flatten the batch and sequence dimensions
+        topk_experts_flat = topk_experts.view(-1, topk_experts.size(-1))
+        topk_probs_flat = topk_probs.view(-1, topk_probs.size(-1))
+        
+        # Initialize layer stats if not present
+        if layer_idx not in self.expert_stats['layer_stats']:
+            self.expert_stats['layer_stats'][layer_idx] = {
+                'total_tokens': 0,
+                'regular_expert_counts': [0] * num_regular_experts,
+                'small_expert_counts': [0] * num_small_experts if num_small_experts > 0 else None,
+                'regular_expert_load': [0.0] * num_regular_experts,
+                'small_expert_load': [0.0] * num_small_experts if num_small_experts > 0 else None
+            }
+        
+        layer_stats = self.expert_stats['layer_stats'][layer_idx]
+        num_tokens = topk_experts_flat.size(0)
+        
+        # Update global stats
+        self.expert_stats['total_tokens'] += num_tokens
+        
+        # Update layer stats
+        layer_stats['total_tokens'] += num_tokens
+        
+        # Track regular experts
+        for expert_idx in range(num_regular_experts):
+            mask = (topk_experts_flat == expert_idx)
+            count = mask.sum().item()
+            load = topk_probs_flat[mask].sum().item()
+            
+            layer_stats['regular_expert_counts'][expert_idx] += count
+            layer_stats['regular_expert_load'][expert_idx] += load
+            
+            if expert_idx not in self.expert_stats['regular_expert_usage']:
+                self.expert_stats['regular_expert_usage'][expert_idx] = 0
+            self.expert_stats['regular_expert_usage'][expert_idx] += count
+        
+        # Track small experts if they exist
+        if num_small_experts > 0:
+            for expert_idx in range(num_small_experts):
+                small_expert_num = expert_idx + num_regular_experts
+                mask = (topk_experts_flat == small_expert_num)
+                count = mask.sum().item()
+                load = topk_probs_flat[mask].sum().item()
+                
+                layer_stats['small_expert_counts'][expert_idx] += count
+                layer_stats['small_expert_load'][expert_idx] += load
+                
+                if expert_idx not in self.expert_stats['small_expert_usage']:
+                    self.expert_stats['small_expert_usage'][expert_idx] = 0
+                self.expert_stats['small_expert_usage'][expert_idx] += count
+
 def _update_expert_stats(self, layer_idx: int, topk_experts: torch.Tensor, 
                         topk_probs: torch.Tensor, num_regular_experts: int,
                         num_small_experts: int, batch_size: int, seq_len: int):