# data_loader.py # Data loader for text classification (AG News, IMDB, DBpedia, Yahoo Answers) for federated experiments. # Note: data-agnostic attack setting ā€” no training-time label flipping is performed. import torch import numpy as np from torch.utils.data import Dataset, DataLoader from transformers import AutoTokenizer import pandas as pd import urllib.request import os from typing import List, Dict class NewsDataset(Dataset): """Custom Dataset for text classification (AG News, IMDB, DBpedia, Yahoo Answers, etc.). Pre-tokenizes the entire text list once at construction so that __getitem__ is a pure tensor-index op. With Qwen-style fast tokenizers, this is bit- identical to per-item tokenization (no RNG, deterministic) but eliminates the ~22.5K redundant tokenize calls per FL round that lazy tokenization caused. Memory cost is negligible (10K samples Ɨ 128 tokens Ɨ int64 ā‰ˆ 10 MB). """ def __init__(self, texts, labels, tokenizer, max_length=128, include_target_mask: bool = False): self.texts = texts self.labels = labels self.tokenizer = tokenizer self.max_length = max_length # include_target_mask is preserved in the signature for backwards # compatibility but is unused in the current codebase (grep confirmed). self.include_target_mask = include_target_mask # Empty dataset (e.g. DataManager.get_empty_loader() for data-agnostic # attackers like AugMP): build zero-row tensors directly. Calling the # fast tokenizer on an empty list raises IndexError on newer # transformers (tokens_and_encodings[0][0]), so short-circuit it. if len(texts) == 0: self._input_ids = torch.empty((0, max_length), dtype=torch.long) self._attention_mask = torch.empty((0, max_length), dtype=torch.long) self._labels = torch.empty((0,), dtype=torch.long) return # Batch tokenize once. Same args as the old per-item call, so the # produced input_ids / attention_mask are identical. encoding = tokenizer( [str(t) for t in texts], truncation=True, padding='max_length', max_length=max_length, return_tensors='pt', ) self._input_ids = encoding['input_ids'] # (N, max_length) self._attention_mask = encoding['attention_mask'] # (N, max_length) self._labels = torch.as_tensor(labels, dtype=torch.long) def __len__(self): return self._input_ids.shape[0] def __getitem__(self, idx): return { 'input_ids': self._input_ids[idx], 'attention_mask': self._attention_mask[idx], 'labels': self._labels[idx], } class DataManager: """Manages text classification data for federated experiments. AG News and Yahoo Answers CSVs live under ``data/ag_news/`` and ``data/yahoo_answers/`` (see ``_load_ag_news`` / ``_load_yahoo_answers``). IMDB and DBpedia load from Hugging Face. """ def __init__(self, num_clients, num_attackers, test_seed, dataset_size_limit=None, batch_size=None, test_batch_size=None, model_name: str = "distilbert-base-uncased", max_length: int = 128, dataset: str = "ag_news"): """ Initialize DataManager. Args: num_clients: Number of federated learning clients (required) num_attackers: Number of attacker clients (required) test_seed: Random seed for test sampling (required) dataset_size_limit: Limit dataset size (None = full dataset). For paper reproduction, use None. When set, only limits training set; test set remains full for fair evaluation. batch_size: Batch size for training data loaders (required) test_batch_size: Batch size for test/validation data loaders (required) model_name: Hugging Face model name for tokenizer initialization max_length: Max token length (AG News: 128, IMDB: 256-512, DBpedia: 512, Yahoo Answers: 256) dataset: 'ag_news' | 'imdb' | 'dbpedia' | 'yahoo_answers'. For ``ag_news`` / ``yahoo_answers``, CSVs are read from ``data/ag_news/`` and ``data/yahoo_answers/`` (see ``data_loader.py``). """ if batch_size is None or test_batch_size is None: raise ValueError("batch_size and test_batch_size must be provided via config (see main.py).") self.num_clients = num_clients self.num_attackers = num_attackers self.test_seed = test_seed self.dataset_size_limit = dataset_size_limit self.batch_size = batch_size self.test_batch_size = test_batch_size self.max_length = max_length self.model_name = model_name self.dataset = dataset.lower() # Load tokenizer self.tokenizer = AutoTokenizer.from_pretrained(model_name) # Handle padding for decoder-only models (GPT-style) if self.tokenizer.pad_token is None: if self.tokenizer.eos_token is not None: self.tokenizer.pad_token = self.tokenizer.eos_token print(f" šŸ“ Set pad_token = eos_token ('{self.tokenizer.eos_token}') for {model_name}") else: self.tokenizer.add_special_tokens({'pad_token': '[PAD]'}) print(f" šŸ“ Added new pad_token '[PAD]' for {model_name}") if self.dataset == "imdb": print("Loading IMDB dataset (stanfordnlp/imdb)...") elif self.dataset == "dbpedia": print("Loading DBpedia dataset (fancyzhx/dbpedia_14)...") elif self.dataset == "yahoo_answers": print("Loading Yahoo Answers dataset (yassiracharki/Yahoo_Answers_10_categories_for_NLP)...") else: print("Loading AG News dataset...") self._load_data() def _load_data(self): """Dispatch to dataset-specific loader.""" if self.dataset == "imdb": self._load_imdb() elif self.dataset == "dbpedia": self._load_dbpedia() elif self.dataset == "yahoo_answers": self._load_yahoo_answers() else: self._load_ag_news() def _load_imdb(self): """Load IMDB dataset from Hugging Face (stanfordnlp/imdb).""" try: from datasets import load_dataset except ImportError: raise ImportError("IMDB requires datasets library. Install: pip install datasets") ds = load_dataset("stanfordnlp/imdb") train_data = ds["train"] test_data = ds["test"] self.train_texts = [str(x) for x in train_data["text"]] self.train_labels = list(train_data["label"]) self.test_texts = [str(x) for x in test_data["text"]] self.test_labels = list(test_data["label"]) print(f" šŸ“Š Full IMDB Dataset: Train={len(self.train_texts)}, Test={len(self.test_texts)}") if self.dataset_size_limit is not None and self.dataset_size_limit > 0: rng = np.random.default_rng(42) n_train = min(self.dataset_size_limit, len(self.train_texts)) n_test = min(int(self.dataset_size_limit * 0.15), len(self.test_texts)) idx_train = rng.choice(len(self.train_texts), n_train, replace=False) idx_test = rng.choice(len(self.test_texts), n_test, replace=False) self.train_texts = [self.train_texts[i] for i in idx_train] self.train_labels = [self.train_labels[i] for i in idx_train] self.test_texts = [self.test_texts[i] for i in idx_test] self.test_labels = [self.test_labels[i] for i in idx_test] print(f" āš ļø Using limited size: Train={len(self.train_texts)}, Test={len(self.test_texts)} (test = train_limit Ɨ 0.15)") print(f" āœ… IMDB ready! Train: {len(self.train_texts)}, Test: {len(self.test_texts)}") def _load_dbpedia(self): """Load DBpedia 14 dataset from Hugging Face (fancyzhx/dbpedia_14).""" try: from datasets import load_dataset except ImportError: raise ImportError("DBpedia requires datasets library. Install: pip install datasets") ds = load_dataset("fancyzhx/dbpedia_14") train_data = ds["train"] test_data = ds["test"] # DBpedia has 'title' and 'content' fields; combine them like AG News train_texts_combined = [f"{str(title)} {str(content)}" for title, content in zip(train_data["title"], train_data["content"])] test_texts_combined = [f"{str(title)} {str(content)}" for title, content in zip(test_data["title"], test_data["content"])] self.train_texts = train_texts_combined self.train_labels = list(train_data["label"]) self.test_texts = test_texts_combined self.test_labels = list(test_data["label"]) print(f" šŸ“Š Full DBpedia Dataset: Train={len(self.train_texts)}, Test={len(self.test_texts)}") if self.dataset_size_limit is not None and self.dataset_size_limit > 0: rng = np.random.default_rng(42) n_train = min(self.dataset_size_limit, len(self.train_texts)) n_test = min(int(self.dataset_size_limit * 0.15), len(self.test_texts)) idx_train = rng.choice(len(self.train_texts), n_train, replace=False) idx_test = rng.choice(len(self.test_texts), n_test, replace=False) self.train_texts = [self.train_texts[i] for i in idx_train] self.train_labels = [self.train_labels[i] for i in idx_train] self.test_texts = [self.test_texts[i] for i in idx_test] self.test_labels = [self.test_labels[i] for i in idx_test] print(f" āš ļø Using limited size: Train={len(self.train_texts)}, Test={len(self.test_texts)} (test = train_limit Ɨ 0.15)") print(f" āœ… DBpedia ready! Train: {len(self.train_texts)}, Test: {len(self.test_texts)}") def _load_yahoo_answers(self): """ Load Yahoo Answers 10-category dataset. 1. Read ``data/yahoo_answers/train.csv`` and ``data/yahoo_answers/test.csv`` if both exist. 2. Otherwise download from Hugging Face and cache under ``data/yahoo_answers/``. """ data_dir = os.path.join("data", "yahoo_answers") os.makedirs(data_dir, exist_ok=True) train_file = os.path.join(data_dir, "train.csv") test_file = os.path.join(data_dir, "test.csv") if os.path.exists(train_file) and os.path.exists(test_file): print(f" āœ… Found local data files in {data_dir}/ directory. Loading...") train_df = pd.read_csv(train_file, header=None, names=['label', 'text'], quoting=1) test_df = pd.read_csv(test_file, header=None, names=['label', 'text'], quoting=1) self.train_texts = train_df['text'].fillna('').astype(str).tolist() self.train_labels = [(int(x) - 1) for x in train_df['label']] self.test_texts = test_df['text'].fillna('').astype(str).tolist() self.test_labels = [(int(x) - 1) for x in test_df['label']] else: try: from datasets import load_dataset except ImportError: raise ImportError("Yahoo Answers requires datasets library. Install: pip install datasets") print(" šŸŒ Local Yahoo Answers CSVs not both under data/yahoo_answers/. Downloading from Hugging Face...") ds = load_dataset("yassiracharki/Yahoo_Answers_10_categories_for_NLP") train_data = ds["train"] test_data = ds["test"] cols = train_data.column_names def _get_col(candidates): for c in candidates: if c in cols: return c return None label_col = _get_col(["class_index", "Class Index", "label"]) or cols[0] title_col = _get_col(["question_title", "Question Title"]) or cols[1] content_col = _get_col(["question_content", "Question Content"]) or cols[2] answer_col = _get_col(["best_answer", "Best Answer"]) or (cols[3] if len(cols) > 3 else None) def _combine_text(t, c, a): parts = [str(x or "").strip() for x in [t, c, a] if x is not None] return " ".join(p for p in parts if p) or " " if answer_col: train_texts = [_combine_text(t, c, a) for t, c, a in zip(train_data[title_col], train_data[content_col], train_data[answer_col])] test_texts = [_combine_text(t, c, a) for t, c, a in zip(test_data[title_col], test_data[content_col], test_data[answer_col])] else: train_texts = [_combine_text(t, c, None) for t, c in zip(train_data[title_col], train_data[content_col])] test_texts = [_combine_text(t, c, None) for t, c in zip(test_data[title_col], test_data[content_col])] train_labels_raw = list(train_data[label_col]) test_labels_raw = list(test_data[label_col]) self.train_texts = train_texts self.train_labels = [int(x) - 1 for x in train_labels_raw] self.test_texts = test_texts self.test_labels = [int(x) - 1 for x in test_labels_raw] train_save = pd.DataFrame({'label': [l + 1 for l in self.train_labels], 'text': self.train_texts}) test_save = pd.DataFrame({'label': [l + 1 for l in self.test_labels], 'text': self.test_texts}) train_save.to_csv(train_file, index=False, header=False, quoting=1) test_save.to_csv(test_file, index=False, header=False, quoting=1) print(f" āœ… Saved to {data_dir}/ for future use.") print(f" šŸ“Š Full Yahoo Answers Dataset: Train={len(self.train_texts)}, Test={len(self.test_texts)}") if self.dataset_size_limit is not None and self.dataset_size_limit > 0: rng = np.random.default_rng(42) n_train = min(self.dataset_size_limit, len(self.train_texts)) n_test = min(int(self.dataset_size_limit * 0.15), len(self.test_texts)) idx_train = rng.choice(len(self.train_texts), n_train, replace=False) idx_test = rng.choice(len(self.test_texts), n_test, replace=False) self.train_texts = [self.train_texts[i] for i in idx_train] self.train_labels = [self.train_labels[i] for i in idx_train] self.test_texts = [self.test_texts[i] for i in idx_test] self.test_labels = [self.test_labels[i] for i in idx_test] print(f" āš ļø Using limited size: Train={len(self.train_texts)}, Test={len(self.test_texts)} (test = train_limit Ɨ 0.15)") print(f" āœ… Yahoo Answers ready! Train: {len(self.train_texts)}, Test: {len(self.test_texts)}") def _load_ag_news(self): """ Load AG News from ``data/ag_news/train.csv`` and ``data/ag_news/test.csv``. Expected CSV format (no header): label, title, text (CharCNN / mhjabreel layout). Missing splits are downloaded from GitHub into ``data/ag_news/`` without overwriting any CSV that already exists locally. """ data_dir = os.path.join("data", "ag_news") os.makedirs(data_dir, exist_ok=True) train_path = os.path.join(data_dir, "train.csv") test_path = os.path.join(data_dir, "test.csv") train_url = "https://raw.githubusercontent.com/mhjabreel/CharCnn_Keras/master/data/ag_news_csv/train.csv" test_url = "https://raw.githubusercontent.com/mhjabreel/CharCnn_Keras/master/data/ag_news_csv/test.csv" try: has_train = os.path.exists(train_path) has_test = os.path.exists(test_path) if has_train and has_test: print(f" āœ… Found local AG News files in {data_dir}/. Loading...") train_df = pd.read_csv(train_path, header=None, names=['label', 'title', 'text']) test_df = pd.read_csv(test_path, header=None, names=['label', 'title', 'text']) else: if not has_train and not has_test: print(" šŸŒ No AG News CSVs under data/ag_news/. Downloading train + test from GitHub...") elif not has_train: print(" šŸŒ Missing train.csv under data/ag_news/. Downloading train split only...") else: print(" šŸŒ Missing test.csv under data/ag_news/. Downloading test split only...") if not has_train: print(f" Train source: {train_url}") with urllib.request.urlopen(train_url, timeout=20) as response: train_raw = response.read().decode('utf-8') with open(train_path, 'w', encoding='utf-8') as f: f.write(train_raw) if not has_test: print(f" Test source: {test_url}") with urllib.request.urlopen(test_url, timeout=20) as response: test_raw = response.read().decode('utf-8') with open(test_path, 'w', encoding='utf-8') as f: f.write(test_raw) train_df = pd.read_csv(train_path, header=None, names=['label', 'title', 'text']) test_df = pd.read_csv(test_path, header=None, names=['label', 'title', 'text']) print(f" āœ… AG News CSVs ready under {data_dir}/.") except Exception as e: print(f"\nāŒ CRITICAL ERROR: Data loading failed: {e}") print("šŸ›‘ STRICT MODE: Synthetic data generation is DISABLED to ensure validity.") print(f" Place train.csv and test.csv under {data_dir}/ or ensure network access for download.") raise e # Process Data # Combine title and text train_df['full_text'] = train_df['title'].astype(str) + ' ' + train_df['text'].astype(str) test_df['full_text'] = test_df['title'].astype(str) + ' ' + test_df['text'].astype(str) # Adjust labels 1-4 -> 0-3 train_df['label'] = train_df['label'] - 1 test_df['label'] = test_df['label'] - 1 # Print full dataset size print(f" šŸ“Š Full AG News Dataset: Train={len(train_df)}, Test={len(test_df)}") # Use full dataset by default # AG News full dataset: ~120,000 training samples, ~7,600 test samples # If dataset_size_limit is set, use it for faster experimentation (not recommended for paper reproduction) if hasattr(self, 'dataset_size_limit') and self.dataset_size_limit is not None: if self.dataset_size_limit > 0: print(f" āš ļø WARNING: Using limited dataset size ({self.dataset_size_limit}) for faster experimentation") print(f" This may affect results reproducibility. For paper reproduction, use full dataset.") train_sample = train_df.sample(n=min(self.dataset_size_limit, len(train_df)), random_state=42) test_sample = test_df.sample(n=min(int(self.dataset_size_limit * 0.15), len(test_df)), random_state=42) else: # Use full dataset train_sample = train_df test_sample = test_df else: # Use full dataset (default, per paper) train_sample = train_df test_sample = test_df self.train_texts = train_sample['full_text'].tolist() self.train_labels = train_sample['label'].tolist() self.test_texts = test_sample['full_text'].tolist() self.test_labels = test_sample['label'].tolist() print(f" āœ… Dataset ready! Train: {len(self.train_texts)}, Test: {len(self.test_texts)}") if len(self.train_texts) < len(train_df) or len(self.test_texts) < len(test_df): print(f" āš ļø Note: Using subset of full dataset (Train: {len(self.train_texts)}/{len(train_df)}, " f"Test: {len(self.test_texts)}/{len(test_df)})") else: print(f" āœ… Using FULL AG News dataset (per paper requirements)") def get_empty_loader(self) -> DataLoader: """Return an empty loader for data-agnostic attackers.""" return DataLoader(NewsDataset([], [], self.tokenizer, max_length=self.max_length), batch_size=self.batch_size, shuffle=False) def get_proxy_eval_loader(self, sample_size: int = 128) -> DataLoader: """ Small clean proxy set for the (omniscient) AugMP attacker's F(w'_g) estimation. Fairness (critical): the proxy is drawn from the TRAINING distribution (``self.train_texts``), NEVER from the test set. The attacker must not be able to optimise its malicious direction against the exact samples used to report clean accuracy / PPL, so the proxy pool is disjoint from the evaluation set by construction. This matches the standard FL threat model where a participant holds a small clean dataset (cf. FLTrust's server root set) and keeps the reported metrics honest. The selection is deterministic (dedicated seed) and class-stratified for a balanced loss estimate. Only AugMP attackers call this method, and it does not touch ``train_texts`` or the client partition, so existing (Hallucination / fedavg / baseline) experiments are unaffected. """ if not self.train_texts: return self.get_empty_loader() # Dedicated seed offset so the proxy selection is independent of the # test-sampling / partition RNGs and reproducible across runs. rng = np.random.default_rng(int(self.test_seed) + 90007) n_train = len(self.train_texts) target = min(int(sample_size), n_train) labels_arr = np.asarray(self.train_labels) classes = np.unique(labels_arr) # Class-stratified quota (balanced proxy -> less biased F(w'_g) estimate). per_class = max(1, target // max(1, len(classes))) chosen: List[int] = [] for c in classes: c_idx = np.where(labels_arr == c)[0] take = min(per_class, len(c_idx)) if take > 0: chosen.extend(rng.choice(c_idx, size=take, replace=False).tolist()) # Top up to `target` from the remaining pool if stratified quota fell short. if len(chosen) < target: remaining = np.setdiff1d(np.arange(n_train), np.asarray(chosen, dtype=int)) if len(remaining) > 0: extra = rng.choice( remaining, size=min(target - len(chosen), len(remaining)), replace=False ) chosen.extend(extra.tolist()) chosen = sorted(chosen) proxy_texts = [self.train_texts[i] for i in chosen] proxy_labels = [self.train_labels[i] for i in chosen] dataset = NewsDataset(proxy_texts, proxy_labels, self.tokenizer, max_length=self.max_length) return DataLoader(dataset, batch_size=self.test_batch_size, shuffle=False) def get_test_loader(self) -> DataLoader: """Get clean global test loader""" test_dataset = NewsDataset(self.test_texts, self.test_labels, self.tokenizer, max_length=self.max_length) return DataLoader(test_dataset, batch_size=self.test_batch_size, shuffle=False)