METRICS_AND_PLOTS_GUIDE.md

# Guide to Metrics, Logs, and Plot Generation

This guide explains how to use, parse, and analyze the validation logs, training logs, router utilization statistics, and plots automatically generated by the updated `train.py` script.

---

## 1. Overview of Generated Files

Upon completing the training run, the following files will be created in your working directory and saved under the checkpoint directory:

| File Name | Format | Description |
| :--- | :--- | :--- |
| **`training_metrics.json`** | JSON | List of dictionaries with training losses recorded at every step: `{"step": s, "train_loss": l}`. |
| **`validation_metrics.json`** | JSON | List of dictionaries with validation metrics evaluated at every 100th step: `{"step": s, "val_loss": l, "val_perplexity": p}`. |
| **`router_details.json`** | JSON | Complete mapping of token indices to experts at every 10th step across all 6 layers. |
| **`router_report.txt`** | Text | Summary of total tokens, bincounts of expert selection per token, and the fractions of tokens receiving no/multiple experts. |
| **`validation_perplexity.png`** | PNG | Line plot of validation perplexity vs global step. |
| **`router_utilization.png`** | PNG | Bar chart of the distribution of the number of experts selected per token (0, 1, 2, 3, or 4). |

---

## 2. Reading & Plotting Learning Curves

You can load and visualize the learning curves at any time using the following Python script:

```python
import json
import matplotlib.pyplot as plt

# 1. Load Metrics
with open("training_metrics.json", "r") as f:
    train_data = json.load(f)
with open("validation_metrics.json", "r") as f:
    val_data = json.load(f)

# Extract lists
train_steps = [item["step"] for item in train_data]
train_losses = [item["train_loss"] for item in train_data]

val_steps = [item["step"] for item in val_data]
val_perplexities = [item["val_perplexity"] for item in val_data]

# 2. Plot Training Loss Curve
plt.figure(figsize=(10, 4))
plt.subplot(1, 2, 1)
plt.plot(train_steps, train_losses, color="#ef4444", label="Training Loss")
plt.xlabel("Global Step")
plt.ylabel("Cross Entropy Loss")
plt.title("Training Loss Curve")
plt.grid(True, linestyle="--", alpha=0.6)
plt.legend()

# 3. Plot Validation Perplexity Curve
plt.subplot(1, 2, 2)
plt.plot(val_steps, val_perplexities, marker="o", color="#3b82f6", label="Validation Perplexity")
plt.xlabel("Global Step")
plt.ylabel("Perplexity")
plt.title("Validation Perplexity Curve")
plt.grid(True, linestyle="--", alpha=0.6)
plt.legend()

plt.tight_layout()
plt.savefig("learning_curves.png", dpi=150)
plt.show()
```

---

## 3. Inspecting Token-to-Expert Mappings

The `router_details.json` file contains a detailed list of dictionaries showing which tokens were routed to which expert block:
```json
[
  {
    "step": 0,
    "layer": 0,
    "expert_assignments": [
      [0, 3, 5, 8, ...],   // Indices selected by Expert 1
      [1, 2, 8, 12, ...],  // Indices selected by Expert 2
      [4, 5, 7, 9, ...],   // Indices selected by Expert 3
      [0, 2, 6, 11, ...]   // Indices selected by Expert 4
    ]
  }
]
```

To reconstruct exactly which token (with its text) was routed to which expert block, use the following code:

```python
import json
from transformers import AutoTokenizer

# Load Tokenizer
tokenizer = AutoTokenizer.from_pretrained("./checkpoints/xpertgpt_fresh/main")

# Load routing details
with open("router_details.json", "r") as f:
    routing_records = json.load(f)

# Inspect a specific step and layer record
record = routing_records[0] # Step 0, Layer 0
step = record["step"]
layer = record["layer"]
assignments = record["expert_assignments"] # List of 4 lists of token indices

print(f"--- Gating Assignments at Step {step}, Layer {layer} ---")
for expert_idx, token_indices in enumerate(assignments):
    print(f"\n[Expert {expert_idx + 1}] selected {len(token_indices)} tokens. Examples:")
    # Print the first 5 token index representations
    for idx in token_indices[:5]:
        print(f"  - Token sequence index {idx}")
```

---

## 4. Re-running the Pipeline to Generate Metrics

To execute the model on Lightning AI and auto-generate these logs and plots, execute:

```bash
# 1. Update the local directory
python -c "import urllib.request; [urllib.request.urlretrieve(f'https://huggingface.co/anonymous_user/correct_small_sw_64_16_8_4_norm_residuls_xpert_strcit_small/raw/main/{f}?v=5', f) for f in ['train.py', 'modeling_xpertgpt.py', 'configuration_xpertgpt.py', 'METRICS_AND_PLOTS_GUIDE.md']]; print('Update complete!')"

# 2. Run pretraining (this will save all logs, reports, and PNG plots locally)
export HF_TOKEN=<your_hf_token>
python train.py --epochs 10 --model-name xpertgpt_fresh

# 3. Push everything back to your HF Hub repository
python train.py --upload --upload-repo correct_small_sw_64_16_8_4_norm_residuls_xpert_strcit_small
```