"""Preprocessing: COVID-debiasing, text normalization, merge, balance, and split datasets.""" import re import random import pandas as pd import numpy as np from sklearn.model_selection import train_test_split from config import CONFIG, COVID_TERMS, REPLACEMENTS from data.load_datasets import load_all_datasets def normalize_text(text): """Normalize tweet-style text to remove stylistic shortcuts. This prevents the model from learning formatting patterns (lowercase, missing punctuation, slang) instead of actual conspiratorial meaning. """ if not isinstance(text, str): return "" # Remove URLs text = re.sub(r"https?://\S+", "", text) # Remove @mentions text = re.sub(r"@\w+", "", text) # Remove hashtag symbols but keep the word text = re.sub(r"#(\w+)", r"\1", text) # Remove RT prefix text = re.sub(r"^RT\s+", "", text) # Normalize multiple spaces text = re.sub(r"\s+", " ", text).strip() # Capitalize first letter of each sentence def capitalize_sentences(t): sentences = re.split(r"([.!?]+\s*)", t) result = [] for i, part in enumerate(sentences): if i == 0 or (i > 0 and re.match(r"[.!?]+\s*", sentences[i - 1])): part = part.lstrip() if part: part = part[0].upper() + part[1:] result.append(part) return "".join(result) text = capitalize_sentences(text) # Add period at end if no sentence-ending punctuation if text and text[-1] not in ".!?": text += "." return text def build_debias_pattern(): """Build a compiled regex pattern for all COVID terms. Returns a list of (compiled_pattern, category) tuples sorted by term length (longest first) to avoid partial matches. """ patterns = [] for category, terms in COVID_TERMS.items(): for term in sorted(terms, key=len, reverse=True): pattern = re.compile(re.escape(term), re.IGNORECASE) patterns.append((pattern, category)) return patterns DEBIAS_PATTERNS = build_debias_pattern() def debias_text(text, rng=None): """Replace COVID-specific terms with generic equivalents. Applied stochastically during training to prevent topic-shortcutting. """ if rng is None: rng = random.Random() for pattern, category in DEBIAS_PATTERNS: replacements = REPLACEMENTS[category] replacement = rng.choice(replacements) text = pattern.sub(replacement, text) return text def cap_and_balance(df, max_samples, seed=42): """Cap a DataFrame to max_samples, maintaining class balance.""" if len(df) <= max_samples: return df label_counts = df["label"].value_counts() per_class = max_samples // len(label_counts) balanced = [] for label in label_counts.index: subset = df[df["label"] == label] n = min(len(subset), per_class) balanced.append(subset.sample(n=n, random_state=seed)) return pd.concat(balanced, ignore_index=True) def merge_datasets(datasets_dict, max_supplement=None, seed=42): """Merge all datasets with optional capping of supplementary data. The primary COVID dataset and hard negatives are kept in full; other supplements are capped. """ if max_supplement is None: max_supplement = CONFIG["max_supplement_samples"] # Always keep these in full (small, curated, critical) keep_full = {"covid_conspiracy", "hard_negatives"} kept = [] supplements = [] for name, df in datasets_dict.items(): if len(df) == 0: continue if name in keep_full: kept.append(df) else: supplements.append(df) if supplements: all_supplements = pd.concat(supplements, ignore_index=True) all_supplements = cap_and_balance(all_supplements, max_supplement, seed) else: all_supplements = pd.DataFrame(columns=["text", "label", "source"]) merged = pd.concat(kept + [all_supplements], ignore_index=True) merged = merged.sample(frac=1, random_state=seed).reset_index(drop=True) return merged def split_data(merged_df, seed=42): """Stratified split into train/val/test + separate COVID-only test set. Returns: train_df, val_df, test_df, covid_test_df """ val_ratio = CONFIG["val_ratio"] test_ratio = CONFIG["test_ratio"] covid_test_size = CONFIG["covid_test_size"] # First, carve out a COVID-only test set covid_only = merged_df[merged_df["source"] == "covid_conspiracy"].copy() non_covid = merged_df[merged_df["source"] != "covid_conspiracy"].copy() if len(covid_only) > covid_test_size: covid_test = covid_only.sample(n=covid_test_size, random_state=seed) covid_remaining = covid_only.drop(covid_test.index) else: # If not enough COVID samples, use 20% for COVID test covid_test = covid_only.sample(frac=0.2, random_state=seed) covid_remaining = covid_only.drop(covid_test.index) # Combine remaining COVID + non-COVID for main splits main_data = pd.concat([covid_remaining, non_covid], ignore_index=True) # Stratified split: first split off test, then split remaining into train/val train_val, test = train_test_split( main_data, test_size=test_ratio, stratify=main_data["label"], random_state=seed, ) relative_val_ratio = val_ratio / (1 - test_ratio) train, val = train_test_split( train_val, test_size=relative_val_ratio, stratify=train_val["label"], random_state=seed, ) return ( train.reset_index(drop=True), val.reset_index(drop=True), test.reset_index(drop=True), covid_test.reset_index(drop=True), ) def compute_class_weights(train_df): """Compute inverse-frequency class weights for balanced loss.""" counts = train_df["label"].value_counts().sort_index() total = len(train_df) weights = total / (len(counts) * counts) return weights.values.astype(np.float32) def preprocess_pipeline(): """Run the full preprocessing pipeline. Returns: train_df, val_df, test_df, covid_test_df, class_weights """ datasets_dict = load_all_datasets() # Normalize all text to remove stylistic shortcuts (capitalization, punctuation) # This prevents the model from learning tweet-formatting as a conspiracy signal print("\nNormalizing text across all datasets...") for name, df in datasets_dict.items(): if len(df) > 0: df["text"] = df["text"].apply(normalize_text) # Remove any empty texts after normalization datasets_dict[name] = df[df["text"].str.strip().str.len() > 0].reset_index(drop=True) print(f" {name}: {len(datasets_dict[name])} samples after normalization") print("\nMerging and balancing datasets...") merged = merge_datasets(datasets_dict) print(f"Merged dataset: {len(merged)} samples") print(f"Label distribution:\n{merged['label'].value_counts()}") print(f"Source distribution:\n{merged['source'].value_counts()}") print("\nSplitting data...") train_df, val_df, test_df, covid_test_df = split_data(merged) print(f"Train: {len(train_df)}, Val: {len(val_df)}, Test: {len(test_df)}, COVID-test: {len(covid_test_df)}") class_weights = compute_class_weights(train_df) print(f"Class weights: {class_weights}") # Show debiasing example covid_samples = train_df[train_df["source"] == "covid_conspiracy"].head(3) if len(covid_samples) > 0: print("\nDebiasing examples:") rng = random.Random(42) for _, row in covid_samples.iterrows(): original = row["text"] debiased = debias_text(original, rng) print(f" Original: {original[:100]}...") print(f" Debiased: {debiased[:100]}...") print() return train_df, val_df, test_df, covid_test_df, class_weights if __name__ == "__main__": train_df, val_df, test_df, covid_test_df, class_weights = preprocess_pipeline()