CFLA 1.0.1__py3-none-any.whl

cfla-1.0.1.dist-info/METADATA

@@ -0,0 +1,249 @@
+Metadata-Version: 2.4
+Name: CFLA
+Version: 1.0.1
+Summary: CFLA - a framework for clustered federated learning algorithms
+Author-email: HOUNSI Antoine <antoinehounsi3@gmail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/Totorino02/CFLA
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: torch
+Requires-Dist: numpy
+Requires-Dist: tqdm
+Requires-Dist: declearn
+Requires-Dist: scikit-learn
+Requires-Dist: pandas
+Requires-Dist: matplotlib
+Provides-Extra: energy
+
+# CFLA — Clustered Federated Learning Algorithms
+
+A Python library implementing and benchmarking **Clustered Federated Learning (CFL)** algorithms. It includes reproductions of state-of-the-art methods from the literature alongside **HCFL**, an original algorithm that automatically discovers the number of clusters without requiring it as a hyperparameter.
+
+---
+
+## Overview
+
+**Federated Learning (FL)** enables training machine learning models on decentralized data without sharing it. Each client trains locally and only sends model updates to a central server.
+
+**Clustered Federated Learning** extends FL by grouping similar clients and training one specialized model per cluster. This improves personalization when client data is heterogeneous (non-IID), which is the common case in practice.
+
+CFLA provides a unified framework to implement, run, and compare CFL algorithms on standard benchmarks, with built-in support for energy consumption monitoring.
+
+---
+
+## Algorithms
+
+### Baselines from the literature
+
+| Algorithm | Clustering strategy | Regularization | Reference |
+|-----------|-------------------|----------------|-----------|
+| **FedAvg** | None — single global model | None | McMahan et al., AISTATS 2017 |
+| **FLHC** | Offline — agglomerative on gradient update vectors | None | Briggs et al., 2020 |
+| **FedGroup** | Offline — agglomerative on model params (cosine) | None | Tan et al., 2022 |
+| **FeSEM** | Online — L2 distance to cluster centers (EM) | None | Li et al., 2021 |
+| **CGPFL** | Online — cosine similarity to cluster centers | `(μ/2)‖ω−Ω_k‖²` | Liu et al. |
+| **IFCA** | Online — empirical risk (loss) per cluster | None | Ghosh et al., NeurIPS 2020 |
+| **LCFed** | Online — PCA low-rank projection + cosine | `(μ/2)‖ω−Ω_k‖² + (λ/2)‖φ−Φ‖²` | Zhang et al., ICASSP 2025 |
+
+### HCFL (original contribution)
+
+HCFL automatically discovers the number of clusters $K^*$ via agglomerative hierarchical clustering on client embedding update vectors, then trains cluster-specific models using a FedProx-style objective toward the cluster center:
+
+$$
+\mathcal{L}_i(\omega) = \mathcal{L}(\omega; D_i) + \frac{\mu}{2} \|\omega - \Omega_k\|^2
+$$
+
+Inter-cluster knowledge sharing is handled **server-side** via a scheduled blending of cluster models toward the global embedding Φ:
+
+$$
+\Omega_k^{(t)} \leftarrow (1 - \lambda^{(t)}) \cdot \text{Avg}\left(\{\omega_i : i \in S_t \cap C_k\}\right) + \lambda^{(t)} \cdot \Phi^{(t)}
+$$
+$$
+\lambda^{(t)} = \frac{\lambda_0}{(1 + \alpha \cdot t)^p}
+$$
+
+This decouples two roles: **μ** controls client-to-cluster fidelity (local), **λ(t)** controls inter-cluster sharing (server-side, decaying).
+
+**Training procedure:**
+1. **Pre-training** (`R_pre` rounds) — FedAvg warm-up to build a meaningful global representation
+2. **Cluster discovery** — each client computes δᵢ = φᵢ - φ⁽⁰⁾ after one local step; agglomerative clustering on {δᵢ} determines K* automatically
+3. **CFL phase** (`T` rounds) — per-cluster training with server-side embedding blending; clusters specialize progressively as λ(t) → 0
+
+---
+
+## Results
+
+Evaluated on MNIST, FEMNIST, and AG News with N=50 clients.
+
+| Method | MNIST Acc ± Std | FEMNIST Acc ± Std | AG News Acc ± Std |
+|--------|----------------|-------------------|-------------------|
+| IFCA† | 97.3 ± 0.5 | 11.4 ± 15.5 | 81.5 ± 19.2 |
+| FLHC† | 97.1 ± 0.4 | **35.2 ± 15.1** | 79.1 ± 22.6 |
+| LCFed† | 82.4 ± 10.5 | 6.9 ± 16.3 | 27.4 ± 30.5 |
+| FeSEM | 90.4 ± 12.2 | 12.8 ± 19.2 | 79.4 ± 22.3 |
+| FedGroup | 91.5 ± 2.1 | 5.2 ± 11.2 | 23.0 ± 25.6 |
+| CGPFL | 50.7 ± 7.9 | 12.7 ± 20.5 | 27.9 ± 26.9 |
+| **HCFL (ours)** | **94.3 ± 0.7** | **32.3 ± 16.7** | **82.0 ± 18.3** |
+
+†: requires K a priori. **Bold**: best no-K method. **Bold + best overall**: AG News.
+
+HCFL is the best no-K method on all three benchmarks, and achieves the best overall accuracy on AG News — surpassing oracle-K methods without any prior knowledge of K.
+
+---
+
+## Project Structure
+
+```
+CFLA/
+├── framework/
+│   ├── client/
+│   │   ├── clientbase.py        # Abstract Client base class
+│   │   ├── client_hcfl.py       # HCFL client
+│   │   ├── client_flhc.py       # FLHC client
+│   │   ├── client_lcfed.py      # LCFed client
+│   │   ├── client_fesem.py      # FeSEM client
+│   │   ├── client_cgpfl.py      # CGPFL client
+│   │   └── client_ifca.py       # IFCA client
+│   ├── server/
+│   │   ├── serverbase.py        # Abstract Server base class
+│   │   ├── server_hcfl.py       # HCFL server
+│   │   ├── server_flhc.py       # FLHC server
+│   │   ├── serveur_lcfed.py     # LCFed server
+│   │   ├── server_fesem.py      # FeSEM server
+│   │   ├── server_cgpfl.py      # CGPFL server
+│   │   └── server_ifca.py       # IFCA server
+│   ├── models/
+│   │   ├── computer_vision.py   # LeNet-5, SplitLeNet5, CNN variants
+│   │   └── nlp_models.py        # DistilBERT-based text encoder
+│   └── common/
+│       └── utils.py             # flatten_params, average_state_dict, cosine_sim
+├── datasets/
+│   ├── femnist.py               # FEMNIST loader + Dirichlet partitioning
+│   └── ag_news.py               # AG News loader + Dirichlet partitioning
+├── experiments/
+│   └── scripts/
+│       ├── run_all_mnist.py     # Run all algorithms on MNIST
+│       ├── run_all_cifar-10.py  # Run all algorithms on CIFAR-10
+│       ├── run_all_femnist.py   # Run all algorithms on FEMNIST
+│       ├── run_all_agnews.py    # Run all algorithms on AG News
+│       ├── hcfl.py
+│       ├── lcfed.py
+│       ├── flhc.py
+│       ├── fedgroup.py
+│       ├── fesem.py
+│       ├── cgpfl.py
+│       ├── ifca.py
+│       └── plot_results.py      # Performance plots + comparison table
+├── pyproject.toml
+└── requirements.dev.txt
+```
+
+---
+
+## Installation
+
+```bash
+git clone https://github.com/Totorino02/CFLA.git
+cd CFLA
+pip install -e .
+pip install -r requirements.dev.txt
+```
+
+or
+
+```bash
+pip install cfla
+```
+
+### Energy monitoring (Linux only)
+
+Energy tracking via RAPL (CPU) and NVML (GPU) is disabled by default (`monitor_energy: False`).
+To enable it on Linux:
+
+```bash
+sudo chmod 444 /sys/class/powercap/intel-rapl:*/energy_uj
+sudo chmod 444 /sys/class/powercap/intel-rapl:*:*/energy_uj
+```
+
+Then set `"monitor_energy": True` in the client args of your experiment script.
+
+---
+
+## Quick Start
+
+Run HCFL on MNIST:
+
+```bash
+python -m experiments.scripts.hcfl
+```
+
+Run all algorithms on a specific dataset:
+
+```bash
+python -m experiments.scripts.run_all_mnist
+python -m experiments.scripts.run_all_femnist
+python -m experiments.scripts.run_all_cifar-10
+python -m experiments.scripts.run_all_agnews
+```
+
+### Generate plots and comparison table
+
+```bash
+python -m experiments.scripts.plot_results --results_dir ./RESULTS/my_run --output_dir ./PLOTS/my_run
+```
+
+Produces:
+- `accuracy_curves.png` — mean accuracy per round with ±1σ band
+- `loss_curves.png` — mean loss per round
+- `accuracy_boxplot.png` — per-client accuracy distribution at the final round
+- `convergence_speed.png` — rounds needed to reach a target accuracy
+- `comparison_table.csv` — final metrics summary
+
+---
+
+## Supported Datasets
+
+| Dataset | Classes | Input | Partitioning |
+|---------|---------|-------|--------------|
+| **MNIST** | 10 | 28×28 grayscale | Structured non-IID (disjoint class groups) |
+| **CIFAR-10** | 10 | 32×32 RGB | Structured non-IID |
+| **FEMNIST** | 62 | 28×28 grayscale | Dirichlet(α) |
+| **AG News** | 4 | Text | Dirichlet(α) |
+
+---
+
+## Output Format
+
+Each experiment writes results to a local `RESULTS/` directory (not tracked in git):
+
+```
+result_{algo}_{dataset}_{timestamp}/
+├── server_metrics.csv           # round, mean_acc, std_acc, mean_loss
+└── client_{id}/
+    └── metrics.csv              # round, loss, accuracy_before, accuracy_after, energy_consumed, energy_ratio
+```
+
+---
+
+## References
+
+- McMahan, B. et al. (2017). Communication-efficient learning of deep networks from decentralized data. *AISTATS*. [arXiv:1602.05629](https://arxiv.org/abs/1602.05629)
+- Briggs, C., Fan, Z., & Andras, P. (2020). Federated learning with hierarchical clustering of local updates to improve training on non-IID data. [arXiv:2004.11791](https://arxiv.org/abs/2004.11791)
+- Ghosh, A. et al. (2020). An efficient framework for clustered federated learning. *NeurIPS*. [arXiv:2006.04088](https://arxiv.org/abs/2006.04088)
+- Tan, Y. et al. (2022). Towards personalized federated learning. *IEEE TNNLS*. [arXiv:2103.00710](https://arxiv.org/abs/2103.00710)
+- Li, X. et al. (2021). FeSEM: Federated learning via expectation maximization. *IEEE TPAMI*.
+- Zhang, Y. et al. (2025). LCFed: An efficient clustered federated learning framework for heterogeneous data. *ICASSP*. [arXiv:2501.01850](https://arxiv.org/abs/2501.01850)
+- Zhang, Y. et al. (2015). Character-level convolutional networks for text classification. *NeurIPS*. [arXiv:1509.01626](https://arxiv.org/abs/1509.01626)
+- Caldas, S. et al. (2019). LEAF: A benchmark for federated settings. [arXiv:1812.01097](https://arxiv.org/abs/1812.01097)
+
+---
+
+## Author
+
+Antoine Hounsi — [antoinehounsi3@gmail.com](mailto:antoinehounsi3@gmail.com)
+
+Master's research project, Université de Lille.

cfla-1.0.1.dist-info/RECORD

@@ -0,0 +1,51 @@
+datasets/__init__.py,sha256=lTKdRvLnWlrKZVTXKzxcIpxmM4DrXQ8FZEkpPh0RL4M,111
+datasets/ag_news.py,sha256=qG8DErSyLrqJENOlPTStUSitjb6bSVA7cbFTv2cl7i0,6679
+datasets/dataset_utils.py,sha256=WQZ6Egi4Zbga7Iuxo_O3gOs6zif8uTfvY5FlHvyMK_M,8580
+datasets/emnist.py,sha256=iVvj3z8JDZSIzmpk_XAe-jnjex3wHy5z3Mk2niz1lfc,857
+datasets/femnist.py,sha256=fhat3y3u3NyXW4l7ktj4hRFQQoPH0RXy2UEi9nHE2mk,6760
+experiments/scripts/cgpfl.py,sha256=lO5v5yx-J1Kyz_LKDrgP5j5cA2CMtUZuErsolid-qxM,2205
+experiments/scripts/fedgroup.py,sha256=iCCna7a2TqLETMnlxOMsLe8LlwTUIblb7IRDr40gMe0,2263
+experiments/scripts/fedper.py,sha256=Q5PbD1lmJR8xUQfYRaD1XGRW-Kq1b7SzyxfIUWM3tdM,2179
+experiments/scripts/fesem.py,sha256=EV2WRmxapqmrrIjb57r3cR81Vc6gJLme-eFC1xQg3yI,2182
+experiments/scripts/flhc.py,sha256=c7Qcy6VTeYkfsXf4Dm9oGgFTJBvmFALtybfWgQmIVLA,2288
+experiments/scripts/hcfl.py,sha256=SSTdeQ8yi_F5oZfg5pyYhYyjBx16e0RuFIKDRXGE0n0,2700
+experiments/scripts/ifca.py,sha256=TzMCJuptKyiugHeQuUOTPDBqia_Y_8aw7OH5YGlt5IQ,2186
+experiments/scripts/lcfed.py,sha256=RGGxiWwnQ7-GOmLj2QhKZ4pZ_R55ILflKm2BIPsGmd8,2517
+experiments/scripts/plot_results.py,sha256=7i1bTfHARq4CJEjfbGoL-2NrRyXs5EJVXwl29_P8IKo,13313
+experiments/scripts/run_all_agnews.py,sha256=DWBdObeJ5zD7WcSVhUwyKvfDc6cSz03z7yTLGfAMHfc,9486
+experiments/scripts/run_all_cifar-10.py,sha256=7kccDeMiBZib9y8Lv9rXAYkR6asBijQLZzy8UQDp8i8,6549
+experiments/scripts/run_all_femnist.py,sha256=bk19LrAkgWaaKJjWU-6bYaTg5IM02mCQDY2kCWyGZq8,10079
+experiments/scripts/run_all_mnist.py,sha256=XMczm3GxCws9K6VlahGDI78IYXOBsQ-dVUveEgk7BxY,6547
+framework/client/__init__.py,sha256=u-lHIoBqU4m6P42roML_MVUJ_ODgFyA8AVxaezmFKF8,40
+framework/client/client_cgpfl.py,sha256=3YpN8I8OP2pofnBYH5Dqf-RkqYonX36HFnu-x_kY3oc,6415
+framework/client/client_fedgroup.py,sha256=AUj2KpZ3mvdxGd4xzMzqLuN2abZJ97qq_dScY9IS9Wo,5690
+framework/client/client_fedper.py,sha256=Ma_tSS9zv9t-jwBXARuX5m4WPVOAFMdwDm40TvILqIw,6613
+framework/client/client_fesem.py,sha256=JDVNWLlcWwXXVp681TYnhFucZg-AfES-KTnC2YA4-mM,5883
+framework/client/client_flhc.py,sha256=8qLWnoG-g9uxsY5Mv0uRRVhQST_fdn-0l7mMb_-VQzg,6280
+framework/client/client_hcfl.py,sha256=Z2WI9p_iCxnYfLfkQeMWjq6qrYM9MW908YTEilK98Ew,10012
+framework/client/client_ifca.py,sha256=oa0PJr_wLsKNoMxoio4_qz_E7z3GqwjF951KDlgu3Rs,6793
+framework/client/client_lcfed.py,sha256=iUfaqk0D0m_QqQFkAP3QS_pfpQgU_BkmAP6VCsWSFHo,10985
+framework/client/client_madmtop.py,sha256=cuowkBbOJZkHvJVmg_sOwD-zwIecMW02ymy0t_jpmYA,3781
+framework/client/clientbase.py,sha256=5QqYVx7Dj8DwoTbSIIqzZt4_NS0lDuusnIPk1dOP3sA,4297
+framework/common/__init__.py,sha256=tWIhSZoWTqEo41QfoyarVSd3_6FmVkVvsSebj0-ww3A,20
+framework/common/parameter_tree.py,sha256=LiFwP2CoHepX_aQ-ruXo1YZvPMhbU5BFYP20NObtte8,3318
+framework/common/utils.py,sha256=2Ep_SJ-hq6tmnIrC4hrGve1S77uVgE6y8N5UpRKrynI,1521
+framework/models/computer_vision.py,sha256=8q0MNeskS0A2f34gyuxTHw6g5u6smFUUlf0XKHVf1OY,5305
+framework/models/large_language_models.py,sha256=Ayksv66rqL7pWbJNbJi_qzYpwM6GBcxvYtYBMNbUIZo,1461
+framework/models/nlp_models.py,sha256=0uoNbYVoeBNQ46w4grpfNxagdgj-CGuvv_icUU_HSPM,2427
+framework/server/__init__.py,sha256=u-lfeZAy9weFoanj4lPeO-S7cJ6d9EtTWXX5erLXMLo,40
+framework/server/server_cgpfl.py,sha256=LkZNPWnVNMxJB_BOvZuihtGAxwI0DsrhpbUMHp1oxBE,7116
+framework/server/server_fedavg.py,sha256=0KEwgRhTk2H8FdNChEbsADAviqM_rrzJz84vQOodAeQ,4761
+framework/server/server_fedgroup.py,sha256=Y_zbfwyb7O4yJ0e4u8cnyonklVp4_eN9Eyo4Z_Hjosg,9106
+framework/server/server_fedper.py,sha256=02SAiQJiqi9TIErXgjkgq836h6Lco2mesLOxTkjqh2Y,5571
+framework/server/server_fesem.py,sha256=2Y3Oi15rD6GeM1vRYd_2rdLcpvJ4feYuNgpSmOHAyMg,7429
+framework/server/server_flhc.py,sha256=9M2zCaIhFUCxUfT2PIDdNdIkajg8BUb6oUXxiyU5vR0,10704
+framework/server/server_hcfl.py,sha256=6G_dWfnu2RykXaIFOtAPZwMtOHBTuWu6Gd5pzbzi8RQ,14923
+framework/server/server_ifca.py,sha256=b_SWbIY-dMmYVTK5Dx8famjpZkDI49nbrF1fsn0xm1g,7060
+framework/server/server_madmtop.py,sha256=RUF_2ivYhDL7bp2HZHkOOgplWQwbU4lOUOg7nWzInZg,8700
+framework/server/serverbase.py,sha256=5Fyzj6mxhXRjEFkh69ld5lUNprPZ1Pt4A1MQHFcQJSc,537
+framework/server/serveur_lcfed.py,sha256=0RVaxGGXtWj48mLlx4QVv0ooSnRxoamPTmfrLMKAmqM,12239
+cfla-1.0.1.dist-info/METADATA,sha256=yPMMg_LLbtQA2RWD7DJpzISKgLVD-QC9VTaivbjsu7k,10219
+cfla-1.0.1.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+cfla-1.0.1.dist-info/top_level.txt,sha256=yuKiyKSdUXIksG2Lz_dat59cba7SaDBq0JN_Sq1W32k,31
+cfla-1.0.1.dist-info/RECORD,,

cfla-1.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

cfla-1.0.1.dist-info/top_level.txt

@@ -0,0 +1,3 @@
+datasets
+experiments
+framework

datasets/__init__.py

@@ -0,0 +1,3 @@
+from .dataset_utils import pathological_partition
+
+__all__ = ["dataset_utils", "emnist", "ag_news", "femnist"]

datasets/ag_news.py

@@ -0,0 +1,198 @@
+"""
+AG News dataset utilities for Federated Learning.
+
+AG News: 4 classes (World=0, Sports=1, Business=2, Sci/Tech=3),
+120 000 train samples, 7 600 test samples.
+
+Loading strategy (in order of preference):
+  1. Local CSV cache (data_root/AG_NEWS/{train,test}.csv)
+  2. Direct download via requests (skips torchtext C extension)
+"""
+
+import csv
+import os
+import re
+from collections import Counter
+
+import numpy as np
+import requests
+import torch
+from torch.utils.data import DataLoader, Dataset, Subset
+
+# ---------------------------------------------------------------------------
+# Raw data loading — bypasses torchtext C extension
+# ---------------------------------------------------------------------------
+
+_AG_NEWS_URLS = {
+    "train": "https://raw.githubusercontent.com/mhjabreel/CharCnn_Keras/master/data/ag_news_csv/train.csv",
+    "test": "https://raw.githubusercontent.com/mhjabreel/CharCnn_Keras/master/data/ag_news_csv/test.csv",
+}
+
+
+def _load_ag_news_csv(split: str, data_root: str) -> list:
+    """
+    Load AG News from local CSV cache or download.
+    Returns list of (label_0indexed, text) tuples.
+    CSV format: class_index (1-4), title, description
+    """
+    cache_dir = os.path.join(data_root, "AG_NEWS")
+    os.makedirs(cache_dir, exist_ok=True)
+    local_path = os.path.join(cache_dir, f"{split}.csv")
+
+    if not os.path.exists(local_path):
+        print(f"Downloading AG News {split} set...")
+        resp = requests.get(_AG_NEWS_URLS[split], timeout=60)
+        resp.raise_for_status()
+        with open(local_path, "wb") as f:
+            f.write(resp.content)
+
+    rows = []
+    with open(local_path, newline="", encoding="utf-8") as f:
+        reader = csv.reader(f)
+        for row in reader:
+            if len(row) < 3:
+                continue
+            label = int(row[0]) - 1  # 1-indexed → 0-indexed
+            text = row[1] + " " + row[2]  # title + description
+            rows.append((label, text))
+    return rows
+
+
+# ---------------------------------------------------------------------------
+# Tokenizer & vocabulary
+# ---------------------------------------------------------------------------
+
+
+def _tokenize(text: str) -> list:
+    """Lowercase + strip non-alphanumeric, split on whitespace."""
+    text = text.lower()
+    text = re.sub(r"[^a-z0-9\s]", " ", text)
+    return text.split()
+
+
+def build_vocab(tokenized_texts: list, max_vocab: int = 25000) -> dict:
+    """
+    Build word→id mapping from a list of token lists.
+    PAD=0, UNK=1, then most-frequent words from index 2.
+    """
+    counter = Counter(tok for tokens in tokenized_texts for tok in tokens)
+    most_common = [w for w, _ in counter.most_common(max_vocab - 2)]
+    vocab = {"<pad>": 0, "<unk>": 1}
+    for i, w in enumerate(most_common, start=2):
+        vocab[w] = i
+    return vocab
+
+
+# ---------------------------------------------------------------------------
+# PyTorch Dataset
+# ---------------------------------------------------------------------------
+
+
+class AGNewsDataset(Dataset):
+    """
+    AG News dataset as a PyTorch Dataset.
+
+    Each item is (token_ids: LongTensor[max_len], label: LongTensor scalar).
+    Labels are 0-indexed (original torchtext labels are 1-indexed).
+    """
+
+    def __init__(self, tokens: list, labels: list, vocab: dict, max_len: int = 128):
+        self.tokens = tokens
+        self.labels = labels
+        self.vocab = vocab
+        self.max_len = max_len
+        self._PAD = vocab["<pad>"]
+        self._UNK = vocab["<unk>"]
+
+    def __len__(self):
+        return len(self.labels)
+
+    def __getitem__(self, idx):
+        ids = [self.vocab.get(w, self._UNK) for w in self.tokens[idx][: self.max_len]]
+        ids += [self._PAD] * (self.max_len - len(ids))
+        return (
+            torch.tensor(ids, dtype=torch.long),
+            torch.tensor(self.labels[idx], dtype=torch.long),
+        )
+
+
+# ---------------------------------------------------------------------------
+# Dataset builder — loads AG News, builds vocab, Dirichlet FL partitioning
+# ---------------------------------------------------------------------------
+
+
+def build_client_datasets_agnews(
+    base_seed: int,
+    run: int,
+    n_clients: int = 40,
+    alpha: float = 0.1,
+    max_len: int = 128,
+    max_vocab: int = 25000,
+    min_samples: int = 64,
+    data_root: str = "./data",
+):
+    """
+    Load AG News (CSV download/cache), build vocab, and partition training data
+    into n_clients via Dirichlet(alpha) on 4 classes.
+
+    Returns
+    -------
+    client_datasets : dict[int, Subset]
+    test_loader     : DataLoader
+    num_classes     : int (4)
+    vocab           : dict (word → id, needed to build the model's vocab_size)
+    """
+    num_classes = 4
+
+    # ---- Load raw data (CSV download, no torchtext C extension needed) ----
+    train_raw = _load_ag_news_csv("train", data_root)
+    test_raw = _load_ag_news_csv("test", data_root)
+
+    # ---- Tokenize ----
+    train_tokens = [_tokenize(text) for _, text in train_raw]
+    test_tokens = [_tokenize(text) for _, text in test_raw]
+    train_labels = [lbl for lbl, _ in train_raw]
+    test_labels = [lbl for lbl, _ in test_raw]
+
+    # ---- Build vocabulary from training set only ----
+    vocab = build_vocab(train_tokens, max_vocab=max_vocab)
+
+    # ---- Full datasets ----
+    train_ds = AGNewsDataset(train_tokens, train_labels, vocab, max_len=max_len)
+    test_ds = AGNewsDataset(test_tokens, test_labels, vocab, max_len=max_len)
+
+    # ---- Dirichlet partitioning on 4 classes ----
+    targets = np.array(train_labels)
+    rng = np.random.default_rng(base_seed + run * 997)
+
+    proportions = rng.dirichlet(alpha * np.ones(num_classes), size=n_clients)
+
+    class_indices = []
+    for c in range(num_classes):
+        idx = np.where(targets == c)[0].tolist()
+        rng.shuffle(idx)
+        class_indices.append(idx)
+
+    client_indices = [[] for _ in range(n_clients)]
+    for c in range(num_classes):
+        idx = class_indices[c]
+        n = len(idx)
+        alloc = (proportions[:, c] / proportions[:, c].sum() * n).astype(int)
+        alloc[-1] = n - alloc[:-1].sum()
+        alloc = np.maximum(alloc, 0)
+        offset = 0
+        for k in range(n_clients):
+            end = offset + alloc[k]
+            client_indices[k].extend(idx[offset:end])
+            offset = end
+
+    for k in range(n_clients):
+        rng.shuffle(client_indices[k])
+
+    client_datasets = {}
+    for cid, indices in enumerate(client_indices):
+        if len(indices) >= min_samples:
+            client_datasets[cid] = Subset(train_ds, indices)
+
+    test_loader = DataLoader(test_ds, batch_size=64, shuffle=False)
+    return client_datasets, test_loader, num_classes, vocab