Update modeling_super_linear.py

modeling_super_linear.py

@@ -356,6 +356,8 @@ class SparseNoisyMoE(nn.Module):
         load_balancing_loss = self.calculate_load_balancing_loss(self.gate_outputs, batch_size)
 
         expert_probs = F.softmax(self.gate_outputs, dim=1)
+        expert_probs = expert_probs[1,:]
+        # Plot the expert probabilities
         import matplotlib.pyplot as plt
 
         plt.figure(figsize=(10, 6))
@@ -363,7 +365,16 @@ class SparseNoisyMoE(nn.Module):
         plt.xlabel('Expert Index')
         plt.ylabel('Probability')
         plt.title('Expert Selection Probabilities')
-        plt.xticks(range(len(expert_probs)), [f'Expert {i}' for i in range(len(expert_probs))])
+        
+        # Create more descriptive x-axis labels
+        if hasattr(self, 'experts') and isinstance(getattr(self, 'experts', None), dict):
+            # If experts are stored in a dictionary with meaningful keys
+            expert_names = list(self.experts.keys())
+            plt.xticks(range(len(expert_probs)), expert_names, rotation=45)
+        else:
+            # Default numbering if expert names aren't available
+            plt.xticks(range(len(expert_probs)), [f'Expert {i}' for i in range(len(expert_probs))])
+        
         plt.tight_layout()
         plt.savefig('expert_probabilities.png')
         plt.close()