diff --git a/README.md b/README.md
index ecb52ce..6c014b7 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,146 @@
 
 # Federated Learning in PyTorch
+
 Implementations of various Federated Learning (FL) algorithms in PyTorch, especially for research purposes.
 
+## Federated Text Classification with DistilBART
+
+This repository contains an implementation of Federated Learning with DistilBART for text classification on the 20 Newsgroups dataset. The implementation includes support for non-IID data distribution across clients using Dirichlet distribution.
+
+## Features
+
+- Federated Learning with DistilBART (distilled version of BART)
+- Support for 20 Newsgroups text classification (20 classes)
+- Non-IID data partitioning using Dirichlet distribution
+- Client-side model training with local updates
+- Centralized model aggregation (FedAvg)
+- Comprehensive evaluation metrics (accuracy, F1, precision, recall)
+- Progress tracking with tqdm and Weights & Biases
+- GPU acceleration support
+- Experiment tracking and visualization
+- Model checkpointing and versioning
+
+## Getting Started
+
+### Prerequisites
+
+- Python 3.8+
+- PyTorch 1.12.0+
+- Transformers 4.18.0+
+- scikit-learn
+- tqdm
+- numpy
+- pandas
+- matplotlib
+- Weights & Biases (`wandb`)
+
+### Installation
+
+1. Clone the repository:
+```bash
+git clone https://github.com/yourusername/FED-OPT-BERT.git
+cd FED-OPT-BERT
+```
+
+2. Install the required packages:
+```bash
+pip install -r requirements.txt
+```
+
+3. Log in to Weights & Biases (if you haven't already):
+```bash
+wandb login
+```
+   Follow the instructions to authenticate with your Weights & Biases account. If you don't have an account, you can create one at [wandb.ai](https://wandb.ai).
+
+### Usage
+
+#### Training
+
+To train the federated DistilBART model on the 20 Newsgroups dataset:
+
+```bash
+python train_distilbart_20news.py \
+    --num_clients 3 \
+    --num_rounds 5 \
+    --epochs_per_client 1 \
+    --batch_size 16 \
+    --learning_rate 2e-5 \
+    --max_grad_norm 1.0 \
+    --data_dir "./data/20news" \
+    --model_save_path "./saved_models/distilbart_20news"
+```
+
+#### Arguments
+
+- `--num_clients`: Number of clients in federated learning (default: 3)
+- `--num_rounds`: Number of federated learning rounds (default: 5)
+- `--epochs_per_client`: Number of local training epochs per client (default: 1)
+- `--batch_size`: Training batch size (default: 16)
+- `--learning_rate`: Learning rate for AdamW optimizer (default: 2e-5)
+- `--max_grad_norm`: Maximum gradient norm for gradient clipping (default: 1.0)
+- `--data_dir`: Directory to store/load the dataset (default: "./data/20news")
+- `--model_save_path`: Path to save the trained model (default: "./saved_models/distilbart_20news")
+
+## Experiment Tracking with Weights & Biases
+
+This project uses Weights & Biases (wandb) for experiment tracking, visualization, and model management. Each training run is automatically logged to your wandb account, where you can:
+
+- Track training and validation metrics in real-time
+- Compare different runs and hyperparameters
+- Monitor system resource usage (CPU/GPU/memory)
+- Save and version model checkpoints
+- Visualize model predictions
+
+### Logged Metrics
+
+- **Training Metrics** (per client, per epoch):
+  - Loss
+  - Accuracy
+  - Precision (weighted)
+  - Recall (weighted)
+  - F1 Score (weighted)
+
+- **Validation Metrics** (per round):
+  - Loss
+  - Accuracy
+  - Precision (weighted)
+  - Recall (weighted)
+  - F1 Score (weighted)
+
+### Viewing Results
+
+1. During or after training, visit your [Weights & Biases dashboard](https://wandb.ai/)
+2. Select your project (`federated-distilbart-20news` by default)
+3. Explore the different tabs:
+   - **Charts**: Interactive plots of all metrics
+   - **System**: Resource utilization
+   - **Models**: Saved model checkpoints
+   - **Files**: Logs and artifacts
+
+## Implementation Details
+
+### Model Architecture
+- Based on DistilBART (distilled version of BART) from Hugging Face
+- Custom classification head for 20 Newsgroups classification
+- Tokenizer: DistilBERT tokenizer with a maximum sequence length of 128 tokens
+
+### Training Process
+1. The global model is initialized with pre-trained DistilBART weights
+2. In each federated round:
+   - A subset of clients is selected
+   - Each client trains the model on its local data
+   - Model updates are sent to the server
+   - The server aggregates the updates using FedAvg
+   - The global model is updated with the aggregated weights
+
+### Evaluation
+- Accuracy, Precision, Recall, F1 Score
+- Confusion matrix
+- Per-class metrics
+- Real-time tracking with Weights & Biases
+- Automatic logging of all metrics and model checkpoints
+
 ## Implementation Details
 ### Datasets
 * Supports all image classification datasets in `torchvision.datasets`.
@@ -54,6 +193,93 @@ Implementations of various Federated Learning (FL) algorithms in PyTorch, especi
 ## Example Commands
 * See shell files prepared in `commands` directory.
 
+### Background Dirichlet alpha sweep (nohup, using experiment runner)
+
+Run a short sweep for Dirichlet α ∈ {0.1, 0.5} sequentially in the background, logging to `nohup_alpha_sweep.log`. Results are written under `--output_dir`.
+
+```bash
+nohup bash -lc '
+for a in 0.1 0.5; do
+  WANDB_MODE=offline /mnt/sda1/Projects/jsl/vp_gitlab/FED/FED-OPT-BERT/FED-OPT-BERT-main/.venv/bin/python \
+    tools/run_20news_experiments.py \
+    --min-clients 2 --max-clients 10 --num-rounds 22 \
+    --participation-rate 1.0 --dirichlet-alpha "$a" --dirichlet-min-size 50 \
+    --output_dir results_distilbart_fed_runs_20news
+done
+' > nohup_alpha_sweep.log 2>&1 &
+```
+
+Notes:
+- `tools/run_20news_experiments.py` forwards flags to `train_distilbart_20news.py`.
+- Omit `--output_dir` to use the default: `results_distilbart_fed_runs_20news`.
+
+## Experiment Results
+
+### Latest Training Run (2025-03-08)
+- **Model**: DistilBART-base
+- **Dataset**: 20 Newsgroups
+- **Configuration**:
+  - Number of clients: 10
+  - Federated rounds: 22
+  - Epochs per client: 1
+  - Batch size: 16
+  - Learning rate: 2e-5
+  - Max sequence length: 128 tokens
+
+### Performance Metrics (Final Round)
+| Metric | Training | Validation |
+|--------|----------|------------|
+| Loss   | 0.644    | 0.062      |
+| Accuracy | 0.798  | 0.724      |
+| Precision | 0.835 | 0.727     |
+| Recall | 0.830    | 0.724      |
+| F1 Score | 0.829  | 0.720      |
+
+### Performance Trends
+- The model shows consistent improvement over federated rounds
+- Training metrics show good convergence
+- Validation metrics indicate the model generalizes well
+- The gap between training and validation metrics suggests some overfitting, which is expected with local training
+
+## Performance Optimization
+
+### Class Imbalance
+- The 20 Newsgroups dataset has relatively balanced classes
+- Consider implementing class weights if needed for specific non-IID scenarios
+
+### Hyperparameter Tuning
+- Experiment with different learning rates and scheduling strategies
+- Try different batch sizes based on available GPU memory
+- Adjust the number of local epochs and federated rounds
+- Use Weights & Biases Sweeps for automated hyperparameter optimization
+
+### Memory Management
+- Gradient accumulation for large batch sizes
+- Mixed precision training (FP16) support
+- Gradient checkpointing for memory efficiency
+
+## Future Work
+
+- [ ] Implement learning rate scheduling with warmup
+- [ ] Add support for more text classification datasets
+- [ ] Implement model compression techniques for edge deployment
+- [ ] Add support for cross-silo federated learning
+- [ ] Add support for federated learning with differential privacy
+- [ ] Implement model distillation for better client-side efficiency
+- [ ] Add support for federated learning with secure aggregation
+
+## Acknowledgements
+
+- [HuggingFace Transformers](https://github.com/huggingface/transformers)
+- [PyTorch](https://pytorch.org/)
+- [scikit-learn](https://scikit-learn.org/)
+- [Weights & Biases](https://wandb.ai/)
+- [FedML](https://fedml.ai/) for federated learning inspiration
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
+
 ## TODO
 - [ ] Support another model, especially lightweight ones for cross-device FL setting. (e.g., [`EdgeNeXt`](https://github.com/mmaaz60/EdgeNeXt))
 - [ ] Support another structured dataset including temporal and tabular data, along with datasets suitable for cross-silo FL setting. (e.g., [`MedMNIST`](https://github.com/MedMNIST/MedMNIST))
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..eeef431
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,8 @@
+torch>=1.9.0
+transformers>=4.11.0
+numpy>=1.20.0
+scikit-learn>=0.24.2
+tqdm>=4.62.0
+pandas>=1.3.0
+matplotlib>=3.4.0
+tensorboard>=2.6.0
diff --git a/results_train_distilbart_20news/plots_train_distilbart_20news/accuracy_over_rounds.png b/results_train_distilbart_20news/plots_train_distilbart_20news/accuracy_over_rounds.png
new file mode 100644
index 0000000..70668ad
Binary files /dev/null and b/results_train_distilbart_20news/plots_train_distilbart_20news/accuracy_over_rounds.png differ
diff --git a/results_train_distilbart_20news/plots_train_distilbart_20news/all_metrics.png b/results_train_distilbart_20news/plots_train_distilbart_20news/all_metrics.png
new file mode 100644
index 0000000..4112106
Binary files /dev/null and b/results_train_distilbart_20news/plots_train_distilbart_20news/all_metrics.png differ
diff --git a/results_train_distilbart_20news/plots_train_distilbart_20news/client_accuracy.png b/results_train_distilbart_20news/plots_train_distilbart_20news/client_accuracy.png
new file mode 100644
index 0000000..62fc2c2
Binary files /dev/null and b/results_train_distilbart_20news/plots_train_distilbart_20news/client_accuracy.png differ
diff --git a/results_train_distilbart_20news/plots_train_distilbart_20news/client_f1.png b/results_train_distilbart_20news/plots_train_distilbart_20news/client_f1.png
new file mode 100644
index 0000000..eaae101
Binary files /dev/null and b/results_train_distilbart_20news/plots_train_distilbart_20news/client_f1.png differ
diff --git a/results_train_distilbart_20news/plots_train_distilbart_20news/client_loss.png b/results_train_distilbart_20news/plots_train_distilbart_20news/client_loss.png
new file mode 100644
index 0000000..88e74fb
Binary files /dev/null and b/results_train_distilbart_20news/plots_train_distilbart_20news/client_loss.png differ
diff --git a/results_train_distilbart_20news/plots_train_distilbart_20news/client_precision.png b/results_train_distilbart_20news/plots_train_distilbart_20news/client_precision.png
new file mode 100644
index 0000000..ab121b9
Binary files /dev/null and b/results_train_distilbart_20news/plots_train_distilbart_20news/client_precision.png differ
diff --git a/results_train_distilbart_20news/plots_train_distilbart_20news/client_recall.png b/results_train_distilbart_20news/plots_train_distilbart_20news/client_recall.png
new file mode 100644
index 0000000..58de761
Binary files /dev/null and b/results_train_distilbart_20news/plots_train_distilbart_20news/client_recall.png differ
diff --git a/results_train_distilbart_20news/plots_train_distilbart_20news/f1_over_rounds.png b/results_train_distilbart_20news/plots_train_distilbart_20news/f1_over_rounds.png
new file mode 100644
index 0000000..3dad8fb
Binary files /dev/null and b/results_train_distilbart_20news/plots_train_distilbart_20news/f1_over_rounds.png differ
diff --git a/results_train_distilbart_20news/plots_train_distilbart_20news/loss_over_rounds.png b/results_train_distilbart_20news/plots_train_distilbart_20news/loss_over_rounds.png
new file mode 100644
index 0000000..3c17637
Binary files /dev/null and b/results_train_distilbart_20news/plots_train_distilbart_20news/loss_over_rounds.png differ
diff --git a/results_train_distilbart_20news/plots_train_distilbart_20news/metrics_correlation.png b/results_train_distilbart_20news/plots_train_distilbart_20news/metrics_correlation.png
new file mode 100644
index 0000000..6dff5be
Binary files /dev/null and b/results_train_distilbart_20news/plots_train_distilbart_20news/metrics_correlation.png differ
diff --git a/results_train_distilbart_20news/plots_train_distilbart_20news/precision_over_rounds.png b/results_train_distilbart_20news/plots_train_distilbart_20news/precision_over_rounds.png
new file mode 100644
index 0000000..fd75036
Binary files /dev/null and b/results_train_distilbart_20news/plots_train_distilbart_20news/precision_over_rounds.png differ
diff --git a/results_train_distilbart_20news/plots_train_distilbart_20news/recall_over_rounds.png b/results_train_distilbart_20news/plots_train_distilbart_20news/recall_over_rounds.png
new file mode 100644
index 0000000..bea5316
Binary files /dev/null and b/results_train_distilbart_20news/plots_train_distilbart_20news/recall_over_rounds.png differ
diff --git a/results_train_distilbart_20news/train_distilbart_20news.csv b/results_train_distilbart_20news/train_distilbart_20news.csv
new file mode 100644
index 0000000..90ec3d9
--- /dev/null
+++ b/results_train_distilbart_20news/train_distilbart_20news.csv
@@ -0,0 +1,133 @@
+round,client_id,epoch,phase,loss,accuracy,precision,recall,f1
+1,1,1,train,2.6849051902168677,0.26724137931034486,0.06194005174690711,0.04774535809018567,0.010467580606195853
+1,7,1,train,2.715851964448628,0.25746516257465163,0.0020501847929476685,0.03185136031851361,0.003358874687966627
+1,6,1,train,2.6962820806001364,0.26542800265428,0.026917830955140686,0.0371599203715992,0.006592986935372853
+1,9,1,train,2.6960480828034252,0.25082946250829463,0.029807081713572978,0.0338420703384207,0.005607975787787075
+1,5,1,train,2.735875870052137,0.23806366047745356,0.00906161380176556,0.04045092838196287,0.00903927298034573
+1,-1,1,validation,0.14352948427411122,0.45977011494252873,0.48376432233643774,0.45977011494252873,0.42754076672257846
+2,2,1,train,2.0178837688345657,0.47679045092838196,0.495915832940753,0.46286472148541113,0.43197510162450464
+2,0,1,train,2.0447773117768135,0.46419098143236076,0.5051981114189248,0.4376657824933687,0.40676891313481467
+2,8,1,train,2.0241483261710718,0.47710683477106836,0.5049263410568796,0.46118115461181153,0.43548885398809434
+2,7,1,train,2.003095091016669,0.5063039150630392,0.5376722728361265,0.46980756469807566,0.43337443796375014
+2,6,1,train,1.978991383000424,0.49834107498341074,0.5356053858573112,0.48573324485733244,0.448764484758555
+2,-1,1,validation,0.11035501324529673,0.5605658709106985,0.5709695449574224,0.5605658709106985,0.5267384601511772
+3,9,1,train,1.6220469349309017,0.5700066357000664,0.5734790568196723,0.5792966157929662,0.5443448982862923
+3,2,1,train,1.624393381570515,0.5576923076923077,0.6053374322177648,0.593501326259947,0.5615731395795185
+3,6,1,train,1.5720229945684734,0.5640345056403451,0.6184436917237973,0.6065029860650298,0.5748777317280129
+3,3,1,train,1.6228506088256835,0.5702917771883289,0.5760348677923678,0.5895225464190982,0.5556175103722812
+3,0,1,train,1.6342491300482498,0.5517241379310345,0.6013840495409762,0.5742705570291777,0.5421094195992179
+3,-1,1,validation,0.09208831072486048,0.6056587091069849,0.5827153783926449,0.6056587091069849,0.5807734640159369
+4,4,1,train,1.397468648458782,0.616710875331565,0.625831258401126,0.6379310344827587,0.6137089003160534
+4,7,1,train,1.3714900286574112,0.6224286662242866,0.6318071878932235,0.6403450564034505,0.6172156394044842
+4,1,1,train,1.3925118258124904,0.6246684350132626,0.6379691725996325,0.6485411140583555,0.6325388345141774
+4,3,1,train,1.3765300098218416,0.6187002652519894,0.6293468866748937,0.6492042440318302,0.622069029258655
+4,0,1,train,1.3892609062947725,0.593501326259947,0.6161858257017678,0.6372679045092838,0.6122911850495263
+4,-1,1,validation,0.08126403824825608,0.6312997347480106,0.6200058101785223,0.6312997347480106,0.6159661656427428
+5,2,1,train,1.2561520024349815,0.6372679045092838,0.6621570630952284,0.6671087533156499,0.6531086628006392
+5,5,1,train,1.3479335797460457,0.6067639257294429,0.6021096624106261,0.6147214854111406,0.5960183768046429
+5,0,1,train,1.19021317205931,0.6591511936339522,0.6534616191784652,0.6704244031830239,0.6515345713681194
+5,4,1,train,1.208593551736129,0.6584880636604774,0.6727151075962255,0.6770557029177718,0.6583017272126832
+5,6,1,train,1.1851808554247807,0.6536164565361646,0.6585705340944236,0.6675514266755143,0.6503066258691133
+5,-1,1,validation,0.07517654341369683,0.6436781609195402,0.6279515045002908,0.6436781609195402,0.6319072368981298
+6,4,1,train,1.0788146486407832,0.6883289124668436,0.6974559571750899,0.7068965517241379,0.6968001724407877
+6,8,1,train,1.1602977815427278,0.6668878566688785,0.665107864754596,0.6721964167219642,0.662770103022675
+6,5,1,train,1.2439554879539891,0.6405835543766578,0.6341471382370215,0.6518567639257294,0.638477531454291
+6,2,1,train,1.1249330730814682,0.6710875331564986,0.688029529498161,0.696949602122016,0.6885451279514319
+6,9,1,train,1.134406452429922,0.6662242866622428,0.7208269725220542,0.6894492368944923,0.6799082069749597
+6,-1,1,validation,0.07176434240121951,0.6560565870910698,0.6404848146156853,0.6560565870910698,0.6446234967047518
+7,7,1,train,1.0684176263056304,0.6741871267418712,0.6845003135615007,0.6921035169210352,0.6835662006276868
+7,2,1,train,1.0218162756217153,0.6949602122015915,0.714308433521704,0.7248010610079576,0.7160429863016533
+7,3,1,train,1.0633896724173897,0.6737400530503979,0.6790534647927896,0.6996021220159151,0.6861099898256982
+7,0,1,train,1.0112992845083537,0.6856763925729443,0.7391840390867015,0.7155172413793104,0.7056893159593742
+7,6,1,train,1.0023052968476949,0.6987392169873923,0.6970399494545987,0.7106834771068348,0.6989587976185133
+7,-1,1,validation,0.06982782338173595,0.6587091069849691,0.6452617274342954,0.6587091069849691,0.6479398960737599
+8,1,1,train,1.0336761380496777,0.696949602122016,0.6997884242229214,0.7049071618037135,0.6986339033143678
+8,8,1,train,1.0115992781363035,0.6921035169210351,0.7038274409366401,0.7120106171201062,0.7020774799931222
+8,2,1,train,0.9550919438663282,0.7062334217506632,0.751886819441725,0.7440318302387268,0.7385814274083008
+8,4,1,train,0.9519191469016828,0.7287798408488064,0.7274237265479977,0.7374005305039788,0.7282961146974243
+8,7,1,train,0.9783030164869209,0.7053749170537492,0.7503738116228379,0.7219641672196416,0.7162110101083876
+8,-1,1,validation,0.06764021357743746,0.6710875331564987,0.6601574409547873,0.6710875331564987,0.6604727123002216
+9,4,1,train,0.8704867428854892,0.7453580901856763,0.7831496852450313,0.7745358090185677,0.7683929616143224
+9,1,1,train,0.9210766823668229,0.7314323607427056,0.7385009989593043,0.7400530503978779,0.7353119470882452
+9,5,1,train,1.0902531978331114,0.6750663129973475,0.7008599754937209,0.6976127320954907,0.6877746788365935
+9,7,1,train,0.9082098148371044,0.7305905773059058,0.7886839718556565,0.7631055076310551,0.7569185965751094
+9,0,1,train,0.8937929300885451,0.7068965517241379,0.770751309448254,0.7473474801061007,0.7379850645329006
+9,-1,1,validation,0.06662465533147756,0.6852343059239611,0.6817369389166323,0.6852343059239611,0.67570937979345
+10,1,1,train,0.861043169623927,0.7513262599469496,0.7774864775549999,0.7619363395225465,0.7584638225514481
+10,7,1,train,0.8184623562975933,0.7558062375580624,0.7933105143619738,0.7790311877903119,0.7740579910119216
+10,8,1,train,0.9168020053913719,0.721300597213006,0.7278370774984496,0.7352355673523556,0.7265415155974952
+10,3,1,train,0.9113883856095766,0.7221485411140585,0.7364825043497005,0.7387267904509284,0.7279069109769941
+10,2,1,train,0.8572485468889538,0.7400530503978779,0.7837303592864203,0.7712201591511937,0.765061472308298
+10,-1,1,validation,0.06581565724549096,0.6799292661361627,0.6728755607671765,0.6799292661361627,0.668699514724873
+11,0,1,train,0.8010347096543563,0.7420424403183025,0.7737660024647224,0.7725464190981433,0.7638582263416192
+11,5,1,train,0.9811709974941455,0.6982758620689655,0.7252609827313584,0.71684350132626,0.7076969420849961
+11,7,1,train,0.7515914562501406,0.7850033178500331,0.8024972769182367,0.794293297942933,0.7894979714095494
+11,2,1,train,0.7877754098490665,0.7679045092838196,0.8055547213157596,0.7937665782493368,0.7900868019295361
+11,6,1,train,0.8515517983781664,0.7438619774386197,0.7595184395430279,0.747843397478434,0.737882849836995
+11,-1,1,validation,0.06513740982659405,0.6878868258178603,0.6851538511256688,0.6878868258178603,0.6794442049519018
+12,5,1,train,0.9174047272456319,0.7155172413793103,0.7536215643221889,0.746684350132626,0.7385492591449629
+12,8,1,train,0.854993506481773,0.7405441274054412,0.7688922169867964,0.7710683477106834,0.762676978246265
+12,2,1,train,0.7235277321777845,0.7864721485411141,0.8212802779893045,0.8110079575596817,0.8083167970782922
+12,1,1,train,0.7841345921943063,0.7725464190981433,0.7802779154351696,0.7765251989389921,0.7723644332819113
+12,3,1,train,0.8404627799987793,0.7360742705570291,0.7621609295492502,0.7672413793103449,0.7578423042524581
+12,-1,1,validation,0.06457038812675274,0.6896551724137931,0.6886480079795667,0.6896551724137931,0.6832377472301185
+13,4,1,train,0.7648063698881551,0.7652519893899205,0.7953833273959504,0.7904509283819628,0.7856439907953848
+13,7,1,train,0.6934914904205423,0.7916390179163902,0.8274127620162098,0.8195089581950896,0.8168411554850808
+13,9,1,train,0.8944190157087226,0.7140013271400133,0.7579541722819191,0.747843397478434,0.7424650072897847
+13,6,1,train,0.7889922215750343,0.7558062375580624,0.7823102386073524,0.7816854678168547,0.7735314642799829
+13,3,1,train,0.7644085381376116,0.7513262599469496,0.7879978143314379,0.7877984084880637,0.7803165010360716
+13,-1,1,validation,0.06320652548819912,0.7038019451812555,0.7070273886824349,0.7038019451812555,0.6985034388346392
+14,7,1,train,0.6457086122349689,0.8195089581950896,0.8519589768848601,0.8453881884538819,0.8434854851136254
+14,3,1,train,0.7199042401815715,0.7778514588859415,0.8163191782452066,0.8103448275862069,0.8051448537315743
+14,0,1,train,0.7137791073636005,0.7745358090185677,0.8075579532461061,0.8017241379310345,0.7948794024432336
+14,5,1,train,0.8343719778876556,0.7427055702917772,0.7711042501452466,0.7679045092838196,0.7608169292655458
+14,6,1,train,0.7345641994162609,0.7717319177173192,0.8114253516018268,0.807564698075647,0.800579229931918
+14,-1,1,validation,0.0630771195751497,0.7020335985853228,0.7033223527145304,0.7020335985853228,0.6963705276373345
+15,1,1,train,0.7165737445417203,0.7864721485411141,0.8181786432159742,0.8076923076923077,0.8076911480582891
+15,5,1,train,0.7762545249964061,0.7685676392572944,0.7974825590941002,0.7957559681697612,0.7901645538820402
+15,0,1,train,0.6444439188430183,0.7964190981432361,0.8292644341214556,0.8249336870026526,0.818206842067155
+15,4,1,train,0.7070423400715777,0.7838196286472149,0.8205626860957995,0.8156498673740054,0.8136842468925555
+15,2,1,train,0.6562258590208857,0.803050397877984,0.8444042962671449,0.8328912466843501,0.8317419821145496
+15,-1,1,validation,0.06282303446062258,0.7020335985853228,0.7050092465974599,0.7020335985853228,0.6964253647524071
+16,1,1,train,0.6326491089243638,0.8136604774535808,0.8353505748661144,0.8269230769230769,0.8255430792260691
+16,4,1,train,0.625974469278988,0.803050397877984,0.8423030166959957,0.8355437665782494,0.8328534398654953
+16,6,1,train,0.6668384539453607,0.7956204379562044,0.8295174882999399,0.8195089581950896,0.811819332533711
+16,2,1,train,0.5947929093712254,0.820291777188329,0.8551849665054782,0.8454907161803713,0.8444007867952845
+16,0,1,train,0.6005680321862823,0.8176392572944298,0.8488263392369823,0.8388594164456233,0.8315040826392864
+16,-1,1,validation,0.06341094142245994,0.7055702917771883,0.7061385505592908,0.7055702917771883,0.6996259584904442
+17,6,1,train,0.6071123742743543,0.8128732581287327,0.8539687400606518,0.8487060384870604,0.8430672318703768
+17,7,1,train,0.5792909651994705,0.8327803583278036,0.86356062903487,0.8573324485733245,0.8562143044536916
+17,2,1,train,0.5420035439102273,0.8415119363395226,0.8792382627182654,0.866710875331565,0.8666246519797806
+17,8,1,train,0.7471050876535867,0.769741207697412,0.8005135286635355,0.7989382879893829,0.7935848362764037
+17,5,1,train,0.7101873415081125,0.7751989389920425,0.8122081674491793,0.8076923076923077,0.8025742024996573
+17,-1,1,validation,0.06329754892173434,0.7064544650751547,0.7092397447708738,0.7064544650751547,0.7016678110750886
+18,3,1,train,0.6340830026488555,0.8103448275862069,0.8322402548865296,0.8282493368700266,0.825476360826032
+18,4,1,train,0.5582005083560944,0.8275862068965517,0.8602304256276676,0.8527851458885941,0.8518578884093959
+18,8,1,train,0.6662523268869048,0.7903118779031187,0.831840847843046,0.827471798274718,0.824599050161182
+18,6,1,train,0.5551853480307679,0.8321167883211679,0.8665662309414818,0.8593231585932316,0.8563260276456337
+18,7,1,train,0.5057421581525552,0.8487060384870604,0.8851971488616297,0.8772395487723955,0.8773767955175911
+18,-1,1,validation,0.06384304549238726,0.7082228116710876,0.7078278743736094,0.7082228116710876,0.7017131756882973
+19,0,1,train,0.5493205014028048,0.8295755968169761,0.8699383453104734,0.8587533156498673,0.8532663191681713
+19,2,1,train,0.5052253065924895,0.8481432360742706,0.889553899147023,0.876657824933687,0.8762953318356291
+19,3,1,train,0.5774284638856587,0.8143236074270557,0.858342327662585,0.8507957559681698,0.8471594445594675
+19,8,1,train,0.6157872724297799,0.8088918380889184,0.8506427484403306,0.8467153284671532,0.8400638799233499
+19,1,1,train,0.6003540794316091,0.8149867374005305,0.8521827078791687,0.8408488063660478,0.8405505550138147
+19,-1,1,validation,0.06291019829578046,0.713527851458886,0.7111757489888112,0.713527851458886,0.7072762972829564
+20,1,1,train,0.5306236866469446,0.8342175066312997,0.877962160269583,0.8680371352785146,0.8693459920080951
+20,8,1,train,0.5594142923229619,0.8307896483078965,0.8713188815942339,0.8666224286662243,0.8644875746184131
+20,0,1,train,0.4896060091100241,0.850132625994695,0.8840465854842596,0.8753315649867374,0.8731285978046188
+20,6,1,train,0.509578465317425,0.845388188453882,0.8782755959089525,0.8732581287325812,0.8706322306582309
+20,9,1,train,0.8097779506131222,0.7485069674850696,0.7785425767501011,0.7763769077637691,0.772569815380377
+20,-1,1,validation,0.06256246843451846,0.713527851458886,0.7147707435949726,0.713527851458886,0.7089609498484601
+21,7,1,train,0.4692117649865778,0.8586595885865959,0.8979647916588874,0.8898473788984738,0.8913353490440867
+21,9,1,train,0.7127552552442802,0.7763769077637691,0.8114744440414368,0.8082282680822827,0.805987655579324
+21,8,1,train,0.5054958763875459,0.8520238885202389,0.8947467636059191,0.8905109489051095,0.8892695527906241
+21,4,1,train,0.5264643795396152,0.8362068965517241,0.8824220348466791,0.8720159151193634,0.8734805383769327
+21,6,1,train,0.44920066292152594,0.8599867285998672,0.8914264506216328,0.8865295288652952,0.8854786107317212
+21,-1,1,validation,0.06272009604015148,0.7214854111405835,0.7245453354825729,0.7214854111405835,0.7187918117700534
+22,7,1,train,0.42075830156865873,0.8719309887193099,0.9166778675011668,0.9097544790975448,0.9109881679970075
+22,5,1,train,0.6507438250278171,0.7911140583554377,0.8343687811458846,0.8295755968169761,0.8282413731423023
+22,3,1,train,0.5165959550557953,0.8362068965517241,0.8729191186964856,0.8693633952254642,0.8685601322658794
+22,9,1,train,0.6444088705864391,0.7976111479761114,0.8346467145150405,0.8301260783012607,0.829480945824145
+22,6,1,train,0.4006392680500683,0.8832116788321167,0.9092030439883393,0.9037823490378235,0.9030797945494602
+22,-1,1,validation,0.06219484595784458,0.7241379310344828,0.7269623861730692,0.7241379310344828,0.7202903631437989
diff --git a/src/datasets/news20.py b/src/datasets/news20.py
new file mode 100644
index 0000000..8830d70
--- /dev/null
+++ b/src/datasets/news20.py
@@ -0,0 +1,168 @@
+import os
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader, Subset
+from sklearn.datasets import fetch_20newsgroups
+from sklearn.model_selection import train_test_split
+from sklearn.feature_extraction.text import TfidfVectorizer
+from transformers import DistilBertTokenizerFast
+
+class News20Dataset(Dataset):
+    def __init__(self, texts, labels, tokenizer, max_length=128):
+        self.texts = texts
+        self.labels = labels
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        
+    def __len__(self):
+        return len(self.texts)
+    
+    def __getitem__(self, idx):
+        text = str(self.texts[idx])
+        label = self.labels[idx]
+        
+        encoding = self.tokenizer(
+            text,
+            max_length=self.max_length,
+            padding='max_length',
+            truncation=True,
+            return_tensors='pt'
+        )
+        
+        return {
+            'input_ids': encoding['input_ids'].flatten(),
+            'attention_mask': encoding['attention_mask'].flatten(),
+            'labels': torch.tensor(label, dtype=torch.long)
+        }
+
+def _split_dirichlet_indices(labels, num_clients, alpha, min_size=10, seed=42):
+    """Partition indices into num_clients using class-wise Dirichlet with concentration alpha.
+
+    Ensures each client has at least min_size samples (retries with new seeds if needed).
+    """
+    rs = np.random.RandomState(seed)
+    labels = np.array(labels)
+    num_classes = int(labels.max()) + 1
+    total_indices = np.arange(len(labels))
+
+    def attempt(seed_offset):
+        rs_local = np.random.RandomState(seed + seed_offset)
+        client_indices = [[] for _ in range(num_clients)]
+        for c in range(num_classes):
+            c_idx = total_indices[labels == c]
+            rs_local.shuffle(c_idx)
+            # Dirichlet proportions for this class across clients
+            proportions = rs_local.dirichlet([alpha] * num_clients)
+            # Convert proportions to split sizes
+            sizes = (proportions * len(c_idx)).astype(int)
+            # Adjust to match exact count by adding remainder to largest bins
+            remainder = len(c_idx) - sizes.sum()
+            if remainder > 0:
+                # add the remainder to clients with largest fractional parts
+                frac = proportions * len(c_idx) - sizes
+                for i in np.argsort(-frac)[:remainder]:
+                    sizes[i] += 1
+            # Now split indices
+            start = 0
+            for i, sz in enumerate(sizes):
+                if sz > 0:
+                    client_indices[i].extend(c_idx[start:start+sz])
+                    start += sz
+        # Shuffle per-client and validate min size
+        for i in range(num_clients):
+            rs_local.shuffle(client_indices[i])
+        if min(len(ci) for ci in client_indices) < min_size:
+            return None
+        return [np.array(ci) for ci in client_indices]
+
+    for retry in range(50):  # avoid infinite loops
+        out = attempt(retry)
+        if out is not None:
+            return out
+    # Fallback: even random split if constraints are too strict
+    indices = np.random.RandomState(seed).permutation(len(labels))
+    return np.array_split(indices, num_clients)
+
+
+def load_20newsgroups(data_dir, num_clients=10, test_size=0.2, random_state=42,
+                      dirichlet_alpha=None, dirichlet_min_size=10):
+    """Load 20 Newsgroups dataset and split it into multiple clients.
+    
+    Args:
+        data_dir (str): Directory to store/load the dataset
+        num_clients (int): Number of clients to split the data into
+        test_size (float): Fraction of data to use for testing
+        random_state (int): Random seed for reproducibility
+        
+    Returns:
+        tuple: (train_datasets, test_datasets, num_classes, tokenizer)
+    """
+    # Create directory if it doesn't exist
+    os.makedirs(data_dir, exist_ok=True)
+    
+    # Load 20 Newsgroups dataset
+    newsgroups_train = fetch_20newsgroups(subset='train', remove=('headers', 'footers', 'quotes'))
+    newsgroups_test = fetch_20newsgroups(subset='test', remove=('headers', 'footers', 'quotes'))
+    
+    # Combine train and test for custom split
+    all_texts = np.concatenate([newsgroups_train.data, newsgroups_test.data])
+    all_labels = np.concatenate([newsgroups_train.target, newsgroups_test.target])
+    
+    # Split into train and test
+    train_texts, test_texts, train_labels, test_labels = train_test_split(
+        all_texts, all_labels, test_size=test_size, random_state=random_state, stratify=all_labels
+    )
+    
+    # Initialize tokenizer
+    tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-uncased')
+    
+    # Create datasets
+    train_dataset = News20Dataset(train_texts, train_labels, tokenizer)
+    test_dataset = News20Dataset(test_texts, test_labels, tokenizer)
+    
+    # Split into clients: IID random or Dirichlet non-IID
+    def to_subsets(dataset, splits):
+        return [Subset(dataset, idx) for idx in splits]
+
+    if dirichlet_alpha is None:
+        # IID random split
+        rs = np.random.RandomState(random_state)
+        train_indices = np.array_split(rs.permutation(len(train_dataset)), num_clients)
+        test_indices = np.array_split(rs.permutation(len(test_dataset)), num_clients)
+    else:
+        # Non-IID via Dirichlet on labels
+        train_indices = _split_dirichlet_indices(
+            labels=train_dataset.labels,
+            num_clients=num_clients,
+            alpha=float(dirichlet_alpha),
+            min_size=int(dirichlet_min_size),
+            seed=random_state,
+        )
+        test_indices = _split_dirichlet_indices(
+            labels=test_dataset.labels,
+            num_clients=num_clients,
+            alpha=float(dirichlet_alpha),
+            min_size=max(1, int(dirichlet_min_size/2)),  # allow smaller test shards
+            seed=random_state + 1,
+        )
+
+    train_datasets = to_subsets(train_dataset, train_indices)
+    test_datasets = to_subsets(test_dataset, test_indices)
+    
+    # Get number of classes
+    num_classes = len(np.unique(all_labels))
+    
+    return train_datasets, test_datasets, num_classes, tokenizer
+
+if __name__ == "__main__":
+    # Example usage
+    data_dir = "./data/20newsgroups"
+    train_datasets, test_datasets, num_classes, tokenizer = load_20newsgroups(
+        data_dir, num_clients=10, test_size=0.2, random_state=42
+    )
+    
+    print(f"Number of classes: {num_classes}")
+    print(f"Number of training clients: {len(train_datasets)}")
+    print(f"Number of test clients: {len(test_datasets)}")
+    print(f"Sample training data size for client 0: {len(train_datasets[0])}")
+    print(f"Sample test data size for client 0: {len(test_datasets[0])}")
diff --git a/src/models/distilbart.py b/src/models/distilbart.py
new file mode 100644
index 0000000..6d69994
--- /dev/null
+++ b/src/models/distilbart.py
@@ -0,0 +1,80 @@
+import torch
+from transformers import DistilBertModel, DistilBertConfig, DistilBertForSequenceClassification
+
+class DistilBART(torch.nn.Module):
+    def __init__(self, num_classes, num_embeddings, embedding_size, hidden_size, dropout, use_pt_model, is_seq2seq=False):
+        super(DistilBART, self).__init__()
+        self.is_seq2seq = is_seq2seq
+        
+        if use_pt_model:  # fine-tuning
+            self.model = DistilBertForSequenceClassification.from_pretrained(
+                'distilbert-base-uncased',
+                num_labels=num_classes,
+                output_attentions=False,
+                output_hidden_states=False,
+            )
+            self.num_embeddings = self.model.config.vocab_size
+            self.embedding_size = self.model.config.dim
+            self.num_hiddens = self.model.config.hidden_size
+            self.dropout = self.model.config.dropout
+        else:  # from scratch
+            self.num_classes = num_classes
+            self.num_embeddings = num_embeddings
+            self.embedding_size = embedding_size
+            self.num_hiddens = hidden_size
+            self.dropout = dropout
+            
+            config = DistilBertConfig(
+                vocab_size=self.num_embeddings,
+                dim=self.embedding_size,
+                hidden_dim=4 * self.embedding_size,  # As per original BERT
+                n_layers=6,  # DistilBERT has 6 layers vs BERT's 12
+                n_heads=8,   # 8 attention heads
+                max_position_embeddings=512,
+                attention_dropout=self.dropout,
+                dropout=self.dropout,
+                num_labels=self.num_classes
+            )
+            self.model = DistilBertForSequenceClassification(config)
+
+    def forward(self, input_ids, attention_mask=None, labels=None):
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            labels=labels
+        )
+        
+        return outputs  # Returns (loss, logits) if labels provided, else just logits
+    
+    def get_embeddings(self):
+        """Get the word embedding layer."""
+        return self.model.distilbert.embeddings.word_embeddings
+    
+    def get_classifier(self):
+        """Get the classifier head."""
+        return self.model.classifier
+    
+    def save_pretrained(self, save_directory):
+        """Save the model and tokenizer to a directory."""
+        self.model.save_pretrained(save_directory)
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        """Load a pretrained model from a directory or HF model hub."""
+        model = DistilBertForSequenceClassification.from_pretrained(
+            pretrained_model_name_or_path,
+            **kwargs
+        )
+        
+        # Create a new instance and replace its model with the loaded one
+        instance = cls(
+            num_classes=model.num_labels,
+            num_embeddings=model.config.vocab_size,
+            embedding_size=model.config.dim,
+            hidden_size=model.config.hidden_size,
+            dropout=model.config.dropout,
+            use_pt_model=True,
+            is_seq2seq=False
+        )
+        instance.model = model
+        return instance
diff --git a/tools/run_20news_experiments.py b/tools/run_20news_experiments.py
new file mode 100644
index 0000000..1eab707
--- /dev/null
+++ b/tools/run_20news_experiments.py
@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+
+import argparse
+import subprocess
+import sys
+from pathlib import Path
+import os
+
+
+def parse_args():
+    p = argparse.ArgumentParser(
+        description="Run multiple 20news federated runs by sweeping num_clients"
+    )
+    p.add_argument("--min-clients", type=int, required=True, help="Minimum number of clients (inclusive)")
+    p.add_argument("--max-clients", type=int, required=True, help="Maximum number of clients (inclusive)")
+    p.add_argument("--num-rounds", type=int, required=True, help="Number of federated rounds per run")
+    # Forwardable optional controls (train script supports these now)
+    p.add_argument("--participation-rate", dest="participation_rate", type=float, default=None,
+                   help="Fraction of clients per round (0 < r <= 1), forwarded to train script")
+    p.add_argument("--min-clients-per-round", dest="min_clients_per_round", type=int, default=None,
+                   help="Minimum clients per round, forwarded to train script")
+    p.add_argument("--dirichlet-alpha", dest="dirichlet_alpha", type=float, default=None,
+                   help="Dirichlet concentration for non-IID split; None means IID (forwarded)")
+    p.add_argument("--dirichlet-min-size", dest="dirichlet_min_size", type=int, default=None,
+                   help="Minimum samples per client for Dirichlet split (forwarded)")
+    # Additional placeholders (not forwarded currently)
+    p.add_argument("--local-epochs", type=int, default=1, help="(unused) Local epochs per client per round")
+    p.add_argument("--batch-size", type=int, default=8, help="(unused) Batch size for training")
+    p.add_argument("--seed", type=int, default=42, help="(unused) Random seed")
+    p.add_argument("--wandb-mode", type=str, default=None, choices=["offline", "online"], help="Set WANDB_MODE for runs")
+    p.add_argument("--python", type=str, default=sys.executable, help="Python interpreter to use (defaults to current)")
+    p.add_argument(
+        "--train-script",
+        type=str,
+        default=str(Path(__file__).resolve().parent.parent / "train_distilbart_20news.py"),
+        help="Path to train_distilbart_20news.py",
+    )
+    p.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="Base directory to store run artifacts; forwarded to train script as --output_dir",
+    )
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+    assert args.min_clients >= 1 and args.max_clients >= args.min_clients
+
+    train_script = Path(args.train_script).resolve()
+    if not train_script.exists():
+        raise FileNotFoundError(f"Train script not found: {train_script}")
+
+    env = os.environ.copy()
+    if args.wandb_mode:
+        env["WANDB_MODE"] = args.wandb_mode
+
+    for n in range(args.min_clients, args.max_clients + 1):
+        # Pass only the arguments supported by train_distilbart_20news.py
+        cmd = [
+            args.python,
+            str(train_script),
+            "--num_clients", str(n),
+            "--num_rounds", str(args.num_rounds),
+        ]
+        if args.participation_rate is not None:
+            cmd += ["--participation_rate", str(args.participation_rate)]
+        if args.min_clients_per_round is not None:
+            cmd += ["--min_clients_per_round", str(args.min_clients_per_round)]
+        if args.dirichlet_alpha is not None:
+            cmd += ["--dirichlet_alpha", str(args.dirichlet_alpha)]
+        if args.dirichlet_min_size is not None:
+            cmd += ["--dirichlet_min_size", str(args.dirichlet_min_size)]
+        if args.output_dir is not None:
+            cmd += ["--output_dir", str(args.output_dir)]
+
+        print("==================================================")
+        print(f"Running: num_clients={n}, num_rounds={args.num_rounds}")
+        print("Command:", " ".join(cmd))
+        print("==================================================")
+        proc = subprocess.run(cmd, env=env)
+        if proc.returncode != 0:
+            print(f"Run failed for num_clients={n} (exit {proc.returncode}). Aborting sweep.")
+            sys.exit(proc.returncode)
+
+    print("All runs completed successfully.")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tools/visualize_runs.py b/tools/visualize_runs.py
new file mode 100644
index 0000000..4f643b2
--- /dev/null
+++ b/tools/visualize_runs.py
@@ -0,0 +1,220 @@
+#!/usr/bin/env python3
+import argparse
+import os
+import glob
+import warnings
+from pathlib import Path
+
+import pandas as pd
+import matplotlib
+matplotlib.use('Agg')
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+sns.set(style="whitegrid")
+
+
+def safe_read_csv(path: Path):
+    try:
+        return pd.read_csv(path)
+    except Exception as e:
+        warnings.warn(f"Failed to read {path}: {e}")
+        return None
+
+
+essential_metric_cols = [
+    'round', 'phase', 'loss', 'accuracy', 'precision', 'recall', 'f1'
+]
+
+
+def plot_metrics_over_rounds(df: pd.DataFrame, out_dir: Path, title_prefix: str = ""):
+    if df is None or df.empty:
+        return
+    # Expect aggregated per-round validation rows or raw rows with phase == validation
+    val = df.copy()
+    if 'phase' in val.columns:
+        val = val[val['phase'].str.contains('val', case=False, na=False)]
+    # Try to ensure numeric types
+    for c in ['round', 'loss', 'accuracy', 'precision', 'recall', 'f1']:
+        if c in val.columns:
+            val[c] = pd.to_numeric(val[c], errors='coerce')
+    val = val.dropna(subset=['round'])
+    if val.empty:
+        return
+
+    metrics = [c for c in ['accuracy', 'f1', 'precision', 'recall', 'loss'] if c in val.columns]
+    if not metrics:
+        return
+
+    # Line plot over rounds for each metric
+    for m in metrics:
+        plt.figure(figsize=(7, 4))
+        sns.lineplot(data=val.sort_values('round'), x='round', y=m, marker='o')
+        plt.title(f"{title_prefix}Validation {m} over rounds")
+        plt.tight_layout()
+        out_path = out_dir / f"val_{m}_over_rounds.png"
+        plt.savefig(out_path, dpi=150)
+        print(f"Saved: {out_path}")
+        plt.close()
+
+
+def plot_client_contributions(df: pd.DataFrame, out_dir: Path):
+    if df is None or df.empty:
+        return
+    # Expect columns like: client_id, split, count or similar
+    # Try to infer
+    cols = {c.lower(): c for c in df.columns}
+    client_col = cols.get('client_id') or cols.get('client') or list(df.columns)[0]
+    count_col = cols.get('count') or cols.get('size') or None
+    if count_col is None:
+        # try to compute from any numeric columns
+        numeric_cols = df.select_dtypes(include='number').columns
+        if len(numeric_cols) >= 1:
+            count_col = numeric_cols[0]
+        else:
+            return
+
+    plt.figure(figsize=(8, 4))
+    order = df.groupby(client_col)[count_col].sum().sort_values(ascending=False).index
+    sns.barplot(data=df, x=client_col, y=count_col, order=order, color="#4C72B0")
+    plt.title("Samples per client")
+    plt.tight_layout()
+    out_path = out_dir / "samples_per_client.png"
+    plt.savefig(out_path, dpi=150)
+    print(f"Saved: {out_path}")
+    plt.close()
+
+
+def plot_label_distribution_heatmap(df: pd.DataFrame, out_dir: Path, title: str):
+    if df is None or df.empty:
+        return
+    # Expect wide-format: client rows x class columns or long with client_id, label, count
+    # Try to detect long format
+    cols = {c.lower(): c for c in df.columns}
+    if {'client_id', 'label', 'count'}.issubset(set(cols.keys())):
+        client = cols['client_id']
+        label = cols['label']
+        count = cols['count']
+        pivot = df.pivot_table(index=client, columns=label, values=count, aggfunc='sum', fill_value=0)
+    else:
+        # if first column is client and others are labels
+        pivot = df.set_index(df.columns[0])
+        # keep only numeric
+        pivot = pivot.select_dtypes(include='number')
+
+    if pivot.empty:
+        return
+
+    plt.figure(figsize=(10, max(4, pivot.shape[0] * 0.4)))
+    sns.heatmap(pivot, annot=False, cmap="Blues")
+    plt.title(title)
+    plt.xlabel("Label")
+    plt.ylabel("Client")
+    plt.tight_layout()
+    fname = "label_distribution_train_heatmap.png" if "train" in title.lower() else "label_distribution_test_heatmap.png"
+    out_path = out_dir / fname
+    plt.savefig(out_path, dpi=150)
+    print(f"Saved: {out_path}")
+    plt.close()
+
+
+def plot_participation(df: pd.DataFrame, out_dir: Path):
+    if df is None or df.empty:
+        return
+    # Expect columns: round, client_id, selected (1/0) or similar
+    cols = {c.lower(): c for c in df.columns}
+    round_col = cols.get('round')
+    client_col = cols.get('client_id') or cols.get('client')
+    selected_col = cols.get('selected') or cols.get('participated') or None
+
+    if not (round_col and client_col):
+        return
+
+    if selected_col is None:
+        # Assume presence indicates participation
+        df['__selected__'] = 1
+        selected_col = '__selected__'
+    
+    # Pivot to rounds x clients (0/1)
+    table = df.pivot_table(index=round_col, columns=client_col, values=selected_col, fill_value=0, aggfunc='max')
+    plt.figure(figsize=(10, max(3, table.shape[1] * 0.3)))
+    sns.heatmap(table.T, cmap='Greens', cbar=False)
+    plt.title('Client Participation by Round')
+    plt.xlabel('Round')
+    plt.ylabel('Client')
+    plt.tight_layout()
+    out_path = out_dir / 'participation_heatmap.png'
+    plt.savefig(out_path, dpi=150)
+    print(f"Saved: {out_path}")
+    plt.close()
+
+
+def find_first(path: Path, patterns):
+    if isinstance(patterns, str):
+        patterns = [patterns]
+    for pat in patterns:
+        found = list(path.glob(pat))
+        if found:
+            return found[0]
+    return None
+
+
+def process_run_dir(run_dir: Path):
+    charts_dir = run_dir / 'charts'
+    charts_dir.mkdir(parents=True, exist_ok=True)
+
+    # Try to find metrics csv
+    metrics_csv = find_first(run_dir, ["training_metrics_*.csv", "metrics*.csv"]) 
+    metrics_df = safe_read_csv(metrics_csv) if metrics_csv else None
+
+    # Contributions and label distributions
+    contrib_csv = find_first(run_dir, ["data_contribution.csv", "client_sizes.csv"]) 
+    contrib_df = safe_read_csv(contrib_csv) if contrib_csv else None
+
+    label_train_csv = find_first(run_dir, ["label_distribution_train.csv", "label_dist_train.csv"]) 
+    label_test_csv = find_first(run_dir, ["label_distribution_test.csv", "label_dist_test.csv"]) 
+    label_train_df = safe_read_csv(label_train_csv) if label_train_csv else None
+    label_test_df = safe_read_csv(label_test_csv) if label_test_csv else None
+
+    participation_csv = find_first(run_dir, ["participation.csv"]) 
+    participation_df = safe_read_csv(participation_csv) if participation_csv else None
+
+    # Plotters
+    plot_metrics_over_rounds(metrics_df, charts_dir)
+    plot_client_contributions(contrib_df, charts_dir)
+    plot_label_distribution_heatmap(label_train_df, charts_dir, "Train label distribution")
+    plot_label_distribution_heatmap(label_test_df, charts_dir, "Test label distribution")
+    plot_participation(participation_df, charts_dir)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Visualize federated run outputs")
+    parser.add_argument("--base-dir", type=str, required=True,
+                        help="Base directory that contains clients_X_rounds_Y/<timestamp>/ run folders")
+    args = parser.parse_args()
+
+    base = Path(args.base_dir)
+    if not base.exists():
+        print(f"Base directory not found: {base}")
+        return 1
+
+    # Expect structure: base/clients_*/TIMESTAMP/
+    client_groups = sorted([p for p in base.iterdir() if p.is_dir()])
+    if not client_groups:
+        print(f"No subdirectories found in {base}")
+        return 0
+
+    run_count = 0
+    for group in client_groups:
+        timestamps = sorted([p for p in group.iterdir() if p.is_dir()])
+        if not timestamps:
+            continue
+        for run_dir in timestamps:
+            process_run_dir(run_dir)
+            run_count += 1
+
+    print(f"Visualization complete. Processed {run_count} run(s).")
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/train_distilbart_20news.py b/train_distilbart_20news.py
new file mode 100644
index 0000000..1901710
--- /dev/null
+++ b/train_distilbart_20news.py
@@ -0,0 +1,443 @@
+import os
+import sys
+import argparse
+import torch
+import numpy as np
+import pandas as pd
+from datetime import datetime
+from torch.utils.data import DataLoader
+from torch.optim import AdamW
+from torch.nn import CrossEntropyLoss
+from tqdm import tqdm
+from sklearn.metrics import precision_recall_fscore_support, accuracy_score
+import wandb
+
+# Add parent directory to path
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from src.models.distilbart import DistilBART
+from src.datasets.news20 import load_20newsgroups
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Federated DistilBART on 20 Newsgroups")
+    parser.add_argument("--num_clients", type=int, default=10, help="Number of clients")
+    parser.add_argument("--num_rounds", type=int, default=22, help="Federated rounds")
+    parser.add_argument("--participation_rate", type=float, default=0.5,
+                        help="Fraction of clients selected per round (0 < r <= 1)")
+    parser.add_argument("--min_clients_per_round", type=int, default=1,
+                        help="Minimum clients selected per round (use 2 for small-N)")
+    parser.add_argument("--dirichlet_alpha", type=float, default=None,
+                        help="Dirichlet concentration for non-IID split; None means IID random split")
+    parser.add_argument("--dirichlet_min_size", type=int, default=10,
+                        help="Minimum samples per client in Dirichlet partition")
+    parser.add_argument("--output_dir", type=str, default="results_distilbart_fed_runs_20news",
+                        help="Base directory to store run artifacts (metrics, checkpoints, logs)")
+    return parser.parse_args()
+
+def train(args):
+    # Configuration
+    config = {
+        'num_clients': args.num_clients,
+        'num_rounds': args.num_rounds,
+        'participation_rate': args.participation_rate,
+        'min_clients_per_round': args.min_clients_per_round,
+        'dirichlet_alpha': args.dirichlet_alpha,
+        'dirichlet_min_size': args.dirichlet_min_size,
+        'epochs_per_client': 1,
+        'batch_size': 16,
+        'learning_rate': 2e-5,
+        'max_seq_length': 128,
+        'model_name': 'distilbart-20news',
+        'project_name': 'federated-distilbart-20news',
+        'data_dir': "./data/20newsgroups",
+        'model_save_path': "./saved_models/distilbart_20news",
+        'output_dir': args.output_dir,
+    }
+    
+    # Initialize wandb
+    wandb.init(
+        project=config['project_name'],
+        name=f"fed-{datetime.now().strftime('%Y%m%d-%H%M%S')}",
+        config=config
+    )
+    
+    # Update config with wandb config (useful for hyperparameter sweeps)
+    config = wandb.config
+    
+    # Unpack config
+    num_clients = config['num_clients']
+    num_rounds = config['num_rounds']
+    participation_rate = float(config.get('participation_rate', 0.5))
+    min_clients_per_round = int(config.get('min_clients_per_round', 1))
+    dirichlet_alpha = config.get('dirichlet_alpha', None)
+    dirichlet_min_size = int(config.get('dirichlet_min_size', 10))
+    epochs_per_client = config['epochs_per_client']
+    batch_size = config['batch_size']
+    learning_rate = config['learning_rate']
+    max_seq_length = config['max_seq_length']
+    data_dir = config['data_dir']
+    model_save_path = config['model_save_path']
+    
+    # Create run-specific results directory
+    output_dir = config.get('output_dir', 'fed_runs')
+    os.makedirs(output_dir, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    run_dir = os.path.join(
+        output_dir,
+        f"clients_{num_clients}_rounds_{num_rounds}",
+        timestamp,
+    )
+    os.makedirs(run_dir, exist_ok=True)
+    results_file = os.path.join(run_dir, f"training_metrics_{timestamp}.csv")
+    
+    # Initialize metrics storage
+    metrics_columns = [
+        'round', 'client_id', 'epoch', 'phase',
+        'loss', 'accuracy', 'precision', 'recall', 'f1'
+    ]
+    metrics_data = []
+    
+    # Set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    
+    # Load data
+    print("Loading 20 Newsgroups dataset...")
+    train_datasets, test_datasets, num_classes, tokenizer = load_20newsgroups(
+        data_dir,
+        num_clients=num_clients,
+        test_size=0.2,
+        random_state=42,
+        dirichlet_alpha=dirichlet_alpha,
+        dirichlet_min_size=dirichlet_min_size,
+    )
+    # Compute and save per-client data contributions and label distributions
+    total_train = sum(len(ds) for ds in train_datasets)
+    total_test = sum(len(ds) for ds in test_datasets)
+    # Access base label arrays from underlying datasets
+    base_train_labels = train_datasets[0].dataset.labels
+    base_test_labels = test_datasets[0].dataset.labels
+
+    contrib_rows = []
+    label_rows_train = []
+    label_rows_test = []
+    for cid in range(num_clients):
+        tr_ds = train_datasets[cid]
+        te_ds = test_datasets[cid]
+        tr_idx = tr_ds.indices
+        te_idx = te_ds.indices
+        tr_size = len(tr_ds)
+        te_size = len(te_ds)
+        contrib_rows.append({
+            'client_id': cid,
+            'train_size': tr_size,
+            'train_prop': tr_size / total_train if total_train > 0 else 0.0,
+            'test_size': te_size,
+            'test_prop': te_size / total_test if total_test > 0 else 0.0,
+        })
+        # Label distributions
+        tr_labels = base_train_labels[tr_idx]
+        te_labels = base_test_labels[te_idx]
+        tr_counts = np.bincount(tr_labels, minlength=num_classes)
+        te_counts = np.bincount(te_labels, minlength=num_classes)
+        for cls in range(num_classes):
+            label_rows_train.append({'client_id': cid, 'class_id': cls, 'count': int(tr_counts[cls])})
+            label_rows_test.append({'client_id': cid, 'class_id': cls, 'count': int(te_counts[cls])})
+
+    pd.DataFrame(contrib_rows).to_csv(os.path.join(run_dir, 'data_contribution.csv'), index=False)
+    pd.DataFrame(label_rows_train).to_csv(os.path.join(run_dir, 'label_distribution_train.csv'), index=False)
+    pd.DataFrame(label_rows_test).to_csv(os.path.join(run_dir, 'label_distribution_test.csv'), index=False)
+    
+    # Initialize global model
+    print("Initializing DistilBART model...")
+    model = DistilBART(
+        num_classes=num_classes,
+        num_embeddings=tokenizer.vocab_size,
+        embedding_size=768,  # DistilBERT base dimension
+        hidden_size=768,
+        dropout=0.1,
+        use_pt_model=True,  # Use pre-trained weights
+        is_seq2seq=False
+    ).to(device)
+    
+    # Define loss function and optimizer
+    criterion = CrossEntropyLoss()
+    
+    # Federated training loop
+    print("Starting federated training...")
+    participation_rows = []
+    for round_num in range(1, num_rounds + 1):
+        print(f"\nRound {round_num}/{num_rounds}")
+        
+        # Randomly select clients for this round based on participation settings
+        # Compute number of participants using ceil, enforce bounds
+        k = int(np.ceil(num_clients * participation_rate))
+        k = max(min_clients_per_round, k)
+        k = min(k, num_clients)
+        selected_clients = np.random.choice(
+            num_clients,
+            size=k,
+            replace=False
+        )
+        
+        # Record participation for this round
+        for cid in range(num_clients):
+            participation_rows.append({'round': round_num, 'client_id': cid, 'selected': int(cid in selected_clients)})
+
+        # Client update phase
+        client_models = []
+        client_sizes = []
+        
+        for client_idx in selected_clients:
+            print(f"\nTraining client {client_idx}...")
+            
+            # Create local model copy
+            local_model = DistilBART(
+                num_classes=num_classes,
+                num_embeddings=tokenizer.vocab_size,
+                embedding_size=768,
+                hidden_size=768,
+                dropout=0.1,
+                use_pt_model=True,
+                is_seq2seq=False
+            ).to(device)
+            local_model.load_state_dict(model.state_dict())
+            
+            # Get client data
+            train_dataset = train_datasets[client_idx]
+            train_loader = DataLoader(
+                train_dataset, 
+                batch_size=batch_size, 
+                shuffle=True,
+                num_workers=4,
+                pin_memory=True
+            )
+            
+            # Local training
+            optimizer = AdamW(local_model.parameters(), lr=learning_rate)
+            local_model.train()
+            
+            for epoch in range(epochs_per_client):
+                total_loss = 0
+                correct = 0
+                total = 0
+                
+                for batch in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs_per_client}"):
+                    input_ids = batch['input_ids'].to(device)
+                    attention_mask = batch['attention_mask'].to(device)
+                    labels = batch['labels'].to(device)
+                    
+                    # Forward pass
+                    outputs = local_model(
+                        input_ids=input_ids,
+                        attention_mask=attention_mask,
+                        labels=labels
+                    )
+                    
+                    loss = outputs.loss
+                    logits = outputs.logits
+                    
+                    # Backward pass and optimize
+                    optimizer.zero_grad()
+                    loss.backward()
+                    optimizer.step()
+                    
+                    # Calculate metrics
+                    total_loss += loss.item()
+                    _, predicted = torch.max(logits, 1)
+                    total += labels.size(0)
+                    correct += (predicted == labels).sum().item()
+                
+                # Calculate metrics
+                avg_loss = total_loss / len(train_loader)
+                accuracy = 100 * correct / total
+                
+                # Get predictions and labels for this epoch
+                all_preds = []
+                all_labels = []
+                with torch.no_grad():
+                    for batch in train_loader:
+                        input_ids = batch['input_ids'].to(device)
+                        attention_mask = batch['attention_mask'].to(device)
+                        labels = batch['labels'].to(device)
+                        
+                        outputs = model(
+                            input_ids=input_ids,
+                            attention_mask=attention_mask
+                        )
+                        _, predicted = torch.max(outputs.logits, 1)
+                        all_preds.extend(predicted.cpu().numpy())
+                        all_labels.extend(labels.cpu().numpy())
+                
+                # Calculate precision, recall, f1 for training
+                precision, recall, f1, _ = precision_recall_fscore_support(
+                    all_labels, all_preds, average='weighted', zero_division=0
+                )
+                
+                # Log metrics to wandb
+                wandb.log({
+                    'round': round_num,
+                    'client': client_idx,
+                    'epoch': epoch + 1,
+                    'train/loss': avg_loss,
+                    'train/accuracy': accuracy / 100,
+                    'train/precision': precision,
+                    'train/recall': recall,
+                    'train/f1': f1
+                }, step=round_num)
+                
+                # Store training metrics
+                metrics_data.append({
+                    'round': round_num,
+                    'client_id': client_idx,
+                    'epoch': epoch + 1,
+                    'phase': 'train',
+                    'loss': avg_loss,
+                    'accuracy': accuracy / 100,  # Convert to 0-1 range
+                    'precision': precision,
+                    'recall': recall,
+                    'f1': f1
+                })
+                
+                print(f"Client {client_idx} - Epoch {epoch+1}: "
+                      f"Loss: {avg_loss:.4f}, Acc: {accuracy:.2f}%")
+            
+            # Store updated model and data size
+            client_models.append(local_model.state_dict())
+            client_sizes.append(len(train_dataset))
+        
+        # Aggregate model updates (Federated Averaging)
+        print("\nAggregating model updates...")
+        global_state = model.state_dict()
+        total_size = sum(client_sizes)
+        
+        # Initialize averaged model parameters
+        for key in global_state.keys():
+            global_state[key] = torch.zeros_like(global_state[key])
+            
+            # Weighted average of client models
+            for i, client_state in enumerate(client_models):
+                weight = client_sizes[i] / total_size
+                global_state[key] += weight * client_state[key]
+        
+        # Update global model
+        model.load_state_dict(global_state)
+        
+        # Evaluate global model on test set
+        print("\nEvaluating global model...")
+        model.eval()
+        
+        # Initialize metrics for validation
+        all_preds = []
+        all_labels = []
+        val_loss = 0.0
+        
+        with torch.no_grad():
+            for client_idx in range(num_clients):  # Evaluate on all clients
+                test_loader = DataLoader(
+                    test_datasets[client_idx],
+                    batch_size=batch_size,
+                    shuffle=False,
+                    num_workers=4,
+                    pin_memory=True
+                )
+                
+                for batch in test_loader:
+                    input_ids = batch['input_ids'].to(device)
+                    attention_mask = batch['attention_mask'].to(device)
+                    labels = batch['labels'].to(device)
+                    
+                    outputs = model(
+                        input_ids=input_ids,
+                        attention_mask=attention_mask,
+                        labels=labels
+                    )
+                    
+                    val_loss += outputs.loss.item()
+                    _, predicted = torch.max(outputs.logits, 1)
+                    
+                    all_preds.extend(predicted.cpu().numpy())
+                    all_labels.extend(labels.cpu().numpy())
+        
+        # Calculate validation metrics
+        val_loss /= len(all_preds)  # Average loss per sample
+        accuracy = accuracy_score(all_labels, all_preds)
+        precision, recall, f1, _ = precision_recall_fscore_support(
+            all_labels, all_preds, average='weighted', zero_division=0
+        )
+        
+        # Log validation metrics to wandb
+        wandb.log({
+            'round': round_num,
+            'val/loss': val_loss,
+            'val/accuracy': accuracy,
+            'val/precision': precision,
+            'val/recall': recall,
+            'val/f1': f1
+        }, step=round_num)
+        
+        # Store validation metrics
+        metrics_data.append({
+            'round': round_num,
+            'client_id': -1,  # -1 indicates global model evaluation
+            'epoch': epochs_per_client,  # Last epoch of the round
+            'phase': 'validation',
+            'loss': val_loss,
+            'accuracy': accuracy,
+            'precision': precision,
+            'recall': recall,
+            'f1': f1
+        })
+        
+        print(f"Round {round_num} - Validation: "
+              f"Loss: {val_loss:.4f}, Acc: {accuracy*100:.2f}%, "
+              f"Precision: {precision:.4f}, Recall: {recall:.4f}, F1: {f1:.4f}")
+        
+        # Save metrics to CSV after each round
+        metrics_df = pd.DataFrame(metrics_data, columns=metrics_columns)
+        metrics_df.to_csv(results_file, index=False)
+        print(f"Metrics saved to {results_file}")
+        
+        # Save model checkpoint into run-specific folder
+        model_checkpoint_dir = os.path.join(run_dir, "checkpoints")
+        os.makedirs(model_checkpoint_dir, exist_ok=True)
+        torch.save({
+            'round': round_num,
+            'model_state_dict': model.state_dict(),
+            'accuracy': accuracy,
+            'loss': val_loss,
+            'f1': f1,
+            'precision': precision,
+            'recall': recall
+        }, os.path.join(model_checkpoint_dir, f'model_round_{round_num}.pt'))
+        
+        print(f"Model checkpoint saved to {model_checkpoint_dir}")
+        
+        # Save model checkpoints to wandb
+        checkpoint_path = os.path.join(model_checkpoint_dir, f'model_round_{round_num}.pt')
+        wandb.save(checkpoint_path)
+    
+    # Save final metrics to a CSV file
+    metrics_df = pd.DataFrame(metrics_data)
+    metrics_file = os.path.join(run_dir, f"training_metrics_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv")
+    metrics_df.to_csv(metrics_file, index=False)
+    wandb.save(metrics_file)
+    # Save participation log
+    if participation_rows:
+        pd.DataFrame(participation_rows).to_csv(os.path.join(run_dir, 'participation.csv'), index=False)
+    
+    # Log the final metrics as a table
+    metrics_table = wandb.Table(dataframe=metrics_df)
+    wandb.log({"metrics_table": metrics_table})
+    
+    # Finish the wandb run
+    wandb.finish()
+
+    # Final metrics summary
+    print("\nTraining completed. Final metrics summary:")
+    metrics_df = pd.read_csv(results_file)
+    print(metrics_df.groupby(['phase']).mean(numeric_only=True)[['loss', 'accuracy', 'precision', 'recall', 'f1']])
+
+if __name__ == "__main__":
+    args = parse_args()
+    train(args)
diff --git a/visualize_metrics.py b/visualize_metrics.py
new file mode 100644
index 0000000..9e36e77
--- /dev/null
+++ b/visualize_metrics.py
@@ -0,0 +1,102 @@
+import pandas as pd
+import matplotlib.pyplot as plt
+import seaborn as sns
+import os
+import numpy as np
+
+def plot_metrics(metrics_file):
+    # Load the metrics
+    df = pd.read_csv(metrics_file)
+    
+    # Create output directory for plots
+    os.makedirs("results/plots", exist_ok=True)
+    
+    # Set style
+    sns.set_style("whitegrid")
+    plt.rcParams["figure.figsize"] = (14, 8)
+    
+    # 1. Plot each metric in a separate figure
+    metrics = ['loss', 'accuracy', 'precision', 'recall', 'f1']
+    
+    # Filter and process data
+    phase_round = df[df['phase'].isin(['train', 'validation'])].groupby(['phase', 'round']).mean(numeric_only=True).reset_index()
+    
+    # Plot each metric separately
+    for metric in metrics:
+        plt.figure(figsize=(10, 6))
+        for phase in ['train', 'validation']:
+            phase_data = phase_round[phase_round['phase'] == phase]
+            if not phase_data.empty:
+                plt.plot(phase_data['round'], phase_data[metric], 
+                        marker='o', label=f'{phase.capitalize()}', linewidth=2)
+        
+        plt.title(f'{metric.capitalize()} over Rounds')
+        plt.xlabel('Round')
+        plt.ylabel(metric.capitalize())
+        plt.legend()
+        plt.grid(True)
+        plt.tight_layout()
+        plt.savefig(f'results/plots/{metric}_over_rounds.png')
+        plt.close()
+    
+    # 2. Plot all metrics in a single figure
+    plt.figure(figsize=(14, 10))
+    for i, metric in enumerate(metrics, 1):
+        plt.subplot(3, 2, i)
+        for phase in ['train', 'validation']:
+            phase_data = phase_round[phase_round['phase'] == phase]
+            if not phase_data.empty:
+                plt.plot(phase_data['round'], phase_data[metric], 
+                        marker='o', label=f'{phase.capitalize()}')
+        
+        plt.title(f'{metric.capitalize()}')
+        plt.xlabel('Round')
+        plt.ylabel(metric.capitalize())
+        plt.legend()
+        plt.grid(True)
+    
+    plt.tight_layout()
+    plt.savefig('results/plots/all_metrics.png')
+    plt.close()
+    
+    # 3. Correlation heatmap
+    plt.figure(figsize=(10, 8))
+    corr = df[metrics].corr()
+    sns.heatmap(corr, annot=True, cmap='coolwarm', vmin=-1, vmax=1, fmt=".2f")
+    plt.title('Metrics Correlation Heatmap')
+    plt.tight_layout()
+    plt.savefig('results/plots/metrics_correlation.png')
+    plt.close()
+    
+    # 4. Client-wise metrics if multiple clients exist
+    if 'client_id' in df.columns:
+        train_df = df[df['phase'] == 'train']
+        if len(train_df['client_id'].unique()) > 1:  # Only if we have multiple clients
+            for metric in metrics:
+                plt.figure(figsize=(10, 6))
+                sns.barplot(data=train_df, x='client_id', y=metric, ci='sd')
+                plt.title(f'Average {metric.capitalize()} by Client')
+                plt.xlabel('Client ID')
+                plt.ylabel(metric.capitalize())
+                plt.grid(True, alpha=0.3)
+                plt.tight_layout()
+                plt.savefig(f'results/plots/client_{metric}.png')
+                plt.close()
+    
+    print("Plots saved to results/plots/ directory")
+
+if __name__ == "__main__":
+    import sys
+    if len(sys.argv) > 1:
+        metrics_file = sys.argv[1]
+    else:
+        # Get the most recent metrics file
+        import glob
+        files = glob.glob("results/training_metrics_*.csv")
+        if not files:
+            print("No metrics files found in results/ directory")
+            sys.exit(1)
+        metrics_file = max(files, key=os.path.getmtime)
+    
+    print(f"Visualizing metrics from: {metrics_file}")
+    plot_metrics(metrics_file)
\ No newline at end of file