dataset.py

from datasets import DatasetDict, load_dataset
from transformers import DataCollatorForSeq2Seq
from constants import MAX_INPUT_LENGTH, MAX_TARGET_LENGTH
import re
import unicodedata

# Pre-compile regex patterns for better performance
SPECIAL_CHAR_PATTERN_ZH = re.compile(r'[^\u4e00-\u9fff\w\s]')
SPECIAL_CHAR_PATTERN_EN = re.compile(r'[^\w\s]')
CHINESE_CHAR_PATTERN = re.compile(r'[\u4e00-\u9fff]')
CONTROL_CHAR_PATTERN = re.compile(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\x9f]')
MULTI_SPACE_PATTERN = re.compile(r'\s+')
REPEATED_PATTERN = re.compile(r'(.{3,})\1{3,}')  # Same 3+ char substring repeated 4+ times

# Fullwidth -> halfwidth mapping table (ASCII range fullwidth chars)
FULLWIDTH_TO_HALFWIDTH = str.maketrans(
    {chr(0xFF01 + i): chr(0x21 + i) for i in range(94)}
)

# Data cleaning constants
MIN_LENGTH = 4
MAX_LENGTH = 256
MAX_LENGTH_RATIO = 5.0
MIN_LENGTH_RATIO = 1.0 / MAX_LENGTH_RATIO
MAX_SPECIAL_CHAR_RATIO = 0.3
MAX_DIGIT_RATIO = 0.5
MAX_WHITESPACE_RATIO = 0.5
MIN_CHINESE_RATIO = 0.1  # Chinese side must have at least 10% Chinese chars
MAX_SOURCE_TARGET_OVERLAP = 0.8  # Filter copy-paste noise

def normalize_text(text: str) -> str:
    """
    Normalize text for consistent tokenization:
    - Remove Unicode control characters
    - Fullwidth -> halfwidth conversion (for punctuation/digits/letters)
    - Normalize Unicode (NFC form)
    - Collapse multiple whitespace into single space
    - Strip leading/trailing whitespace
    """
    text = CONTROL_CHAR_PATTERN.sub('', text)
    text = text.translate(FULLWIDTH_TO_HALFWIDTH)
    text = unicodedata.normalize('NFC', text)
    text = MULTI_SPACE_PATTERN.sub(' ', text)
    return text.strip()


def build_dataset() -> DatasetDict:
    print("Loading WMT19 dataset...")
    dataset = load_dataset("wmt19", "zh-en", split="train")

    print("Shuffling dataset for better quality...")
    shuffled_dataset = dataset.shuffle(seed=42)

    # Select training and validation data
    # Increased training data for better performance
    train_dataset = shuffled_dataset.select(range(600000))
    validation_dataset = shuffled_dataset.select(range(600000, 602000))
    test_dataset = load_dataset("wmt19", "zh-en", split="validation")

    return DatasetDict({
        "train": train_dataset,
        "validation": validation_dataset,
        "test": test_dataset
    })

def preprocess_function(examples, tokenizer, max_input_length, max_target_length):
    # Set source and target languages for NLLB
    tokenizer.src_lang = "zho_Hans"
    tokenizer.tgt_lang = "eng_Latn"
    inputs = [normalize_text(ex["zh"]) for ex in examples["translation"]]
    targets = [normalize_text(ex["en"]) for ex in examples["translation"]]
    model_inputs = tokenizer(inputs, max_length=max_input_length, truncation=True)
    labels = tokenizer(text_target=targets, max_length=max_target_length, truncation=True)
    model_inputs["labels"] = labels["input_ids"]
    return model_inputs

def clean_data(example):
    """
    Comprehensive data cleaning function with optimized performance.
    Filters out low-quality translation pairs based on multiple criteria.
    Uses early returns and pre-compiled regex for better performance.
    """
    # Early return for invalid structure (most common failure case)
    if not isinstance(example.get('translation'), dict):
        return False

    translation = example['translation']
    if 'zh' not in translation or 'en' not in translation:
        return False

    # Normalize and strip whitespace
    zh_text = normalize_text(translation['zh'])
    en_text = normalize_text(translation['en'])

    # Early return for empty strings
    zh_len = len(zh_text)
    en_len = len(en_text)
    if zh_len == 0 or en_len == 0:
        return False

    # Length filtering: early return for extreme lengths (fast check)
    if zh_len < MIN_LENGTH or en_len < MIN_LENGTH:
        return False

    if zh_len > MAX_LENGTH or en_len > MAX_LENGTH:
        return False

    # Length ratio filtering: filter misaligned translations
    length_ratio = zh_len / en_len
    if length_ratio > MAX_LENGTH_RATIO or length_ratio < MIN_LENGTH_RATIO:
        return False

    # Language consistency: Chinese side must actually contain Chinese characters
    chinese_chars = len(CHINESE_CHAR_PATTERN.findall(zh_text))
    if chinese_chars / zh_len < MIN_CHINESE_RATIO:
        return False

    # Language consistency: English side should not contain too many Chinese characters
    chinese_in_en = len(CHINESE_CHAR_PATTERN.findall(en_text))
    if en_len > 0 and chinese_in_en / en_len > 0.1:
        return False

    # Source-target overlap: filter copy-paste noise (zh and en almost identical)
    common_chars = set(zh_text) & set(en_text)
    overlap = len(common_chars) / max(len(set(zh_text)), len(set(en_text)), 1)
    if overlap > MAX_SOURCE_TARGET_OVERLAP:
        return False

    # Repeated pattern detection: filter garbled/noisy text
    if REPEATED_PATTERN.search(zh_text) or REPEATED_PATTERN.search(en_text):
        return False

    # Character quality checks (more expensive, done after length checks)
    # Special character ratio
    special_chars_zh = len(SPECIAL_CHAR_PATTERN_ZH.findall(zh_text))
    special_chars_en = len(SPECIAL_CHAR_PATTERN_EN.findall(en_text))

    if special_chars_zh / zh_len > MAX_SPECIAL_CHAR_RATIO:
        return False

    if special_chars_en / en_len > MAX_SPECIAL_CHAR_RATIO:
        return False

    # Digit ratio: filter sentences that are mostly numbers
    digit_count_zh = sum(c.isdigit() for c in zh_text)
    digit_count_en = sum(c.isdigit() for c in en_text)

    if digit_count_zh / zh_len > MAX_DIGIT_RATIO:
        return False

    if digit_count_en / en_len > MAX_DIGIT_RATIO:
        return False

    # Whitespace ratio: filter sentences with excessive whitespace
    whitespace_ratio_zh = sum(c.isspace() for c in zh_text) / zh_len
    whitespace_ratio_en = sum(c.isspace() for c in en_text) / en_len

    if whitespace_ratio_zh > MAX_WHITESPACE_RATIO or whitespace_ratio_en > MAX_WHITESPACE_RATIO:
        return False

    return True


def basic_validation_filter(example):
    """
    Lightweight validation filter that only checks basic structure.
    Used for validation and test sets to maintain data distribution consistency.
    Does NOT perform quality filtering like clean_data.
    """
    if not isinstance(example.get('translation'), dict):
        return False

    translation = example['translation']
    if 'zh' not in translation or 'en' not in translation:
        return False

    # Only check for non-empty strings (basic structure validation)
    zh_text = translation.get('zh', '').strip()
    en_text = translation.get('en', '').strip()

    return len(zh_text) > 0 and len(en_text) > 0


def remove_duplicates(dataset, text_key='translation'):
    """
    Remove duplicate translation pairs based on exact text matching.
    Keeps the first occurrence of each unique pair.
    Note: This uses a closure to maintain the seen set per dataset.
    """
    seen = set()

    def is_unique(example):
        if text_key not in example:
            return True
        translation = example[text_key]
        if not isinstance(translation, dict) or 'zh' not in translation or 'en' not in translation:
            return True

        # Create a unique key from both texts
        pair_key = (translation['zh'].strip(), translation['en'].strip())
        if pair_key in seen:
            return False
        seen.add(pair_key)
        return True

    return dataset.filter(is_unique)


def preprocess_data(raw_datasets: DatasetDict, tokenizer, remove_dups: bool = True) -> DatasetDict:
    """
    Preprocess and clean the datasets.
    Applies comprehensive data cleaning to training set only.
    Validation set only gets basic structure validation (no quality filtering).
    Test set is kept completely unchanged to maintain original distribution for fair evaluation.

    Args:
        raw_datasets: DatasetDict containing train, validation, and test splits
        tokenizer: Tokenizer for encoding text
        remove_dups: Whether to remove duplicate translation pairs from training set (default: True)
    """
    print("Cleaning training data...")

    # Apply comprehensive data cleaning to training set only
    if "train" in raw_datasets:
        initial_train_size = len(raw_datasets["train"])
        raw_datasets["train"] = raw_datasets["train"].filter(clean_data)
        cleaned_train_size = len(raw_datasets["train"])

        # Remove duplicates if enabled
        if remove_dups:
            raw_datasets["train"] = remove_duplicates(raw_datasets["train"])
            after_dedup = len(raw_datasets["train"])
            print(f"  Training set: {initial_train_size} -> {cleaned_train_size} (cleaned) -> {after_dedup} (dedup) samples "
                  f"({(1 - after_dedup/initial_train_size)*100:.2f}% total filtered)")
        else:
            print(f"  Training set: {initial_train_size} -> {cleaned_train_size} samples "
                  f"({(1 - cleaned_train_size/initial_train_size)*100:.2f}% filtered)")

    # Apply only basic structure validation to validation set (no quality filtering)
    if "validation" in raw_datasets:
        initial_val_size = len(raw_datasets["validation"])
        raw_datasets["validation"] = raw_datasets["validation"].filter(basic_validation_filter)
        filtered_val_size = len(raw_datasets["validation"])
        print(f"  Validation set: {initial_val_size} -> {filtered_val_size} samples "
              f"(basic structure validation only, {(1 - filtered_val_size/initial_val_size)*100:.2f}% filtered)")

    # Test set is kept completely unchanged to maintain original distribution for fair evaluation
    if "test" in raw_datasets:
        test_size = len(raw_datasets["test"])
        print(f"  Test set: {test_size} samples (kept unchanged, no filtering applied)")

    print("Tokenizing datasets...")
    return raw_datasets.map(
        function=lambda examples: preprocess_function(examples, tokenizer, MAX_INPUT_LENGTH, MAX_TARGET_LENGTH),
        batched=True,
        remove_columns=raw_datasets["train"].column_names,
    )