| from __future__ import annotations |
|
|
| import re |
| import unicodedata |
|
|
| import pandas as pd |
| from tqdm.auto import tqdm |
|
|
| from .schema import OPTIONAL_COLUMNS, REQUIRED_COLUMNS |
| from .taxonomy import canonicalize_category |
|
|
|
|
| def normalize_text(text: object) -> str: |
| if text is None or pd.isna(text): |
| return "" |
| value = unicodedata.normalize("NFKC", str(text)).strip().lower() |
| value = re.sub(r"\s+", " ", value) |
| return value |
|
|
|
|
| def normalize_category(label: object) -> str: |
| if label is None or pd.isna(label): |
| return "" |
| value = unicodedata.normalize("NFKC", str(label)).strip().lower() |
| |
| |
| |
| value = value.replace("-", " ") |
| return re.sub(r"\s+", " ", value) |
|
|
|
|
| def normalize_country(value: object) -> str: |
| if value is None or pd.isna(value): |
| return "unknown" |
| normalized = str(value).strip().casefold() |
| aliases = { |
| "usa": "USA", |
| "us": "USA", |
| "united states": "USA", |
| "uk": "UK", |
| "united kingdom": "UK", |
| "canada": "Canada", |
| "australia": "Australia", |
| "india": "India", |
| } |
| return aliases.get(normalized, str(value).strip()) or "unknown" |
|
|
|
|
| def normalize_optional_text(value: object) -> str: |
| if value is None or pd.isna(value): |
| return "unknown" |
| normalized = str(value).strip() |
| return normalized or "unknown" |
|
|
|
|
| def prepare_dataframe( |
| frame: pd.DataFrame, |
| *, |
| progress: bool = True, |
| chunk_size: int = 250_000, |
| canonicalize: bool = False, |
| strict_taxonomy: bool = True, |
| ) -> pd.DataFrame: |
| missing = set(REQUIRED_COLUMNS) - set(frame.columns) |
| if missing: |
| raise ValueError(f"Missing required columns: {sorted(missing)}") |
| if chunk_size < 1: |
| raise ValueError("chunk_size must be >= 1") |
|
|
| columns = list(REQUIRED_COLUMNS) + [c for c in OPTIONAL_COLUMNS if c in frame.columns] |
| result = frame.loc[:, columns].copy() |
| iterator = range(0, len(result), chunk_size) |
| total_chunks = (len(result) + chunk_size - 1) // chunk_size |
| progress_bar = tqdm( |
| iterator, |
| total=total_chunks, |
| unit="chunks", |
| desc="Normalizing", |
| disable=not progress, |
| ) |
| chunks: list[pd.DataFrame] = [] |
| unknown_labels: set[tuple[str, str]] = set() |
|
|
| for start in progress_bar: |
| chunk = result.iloc[start : start + chunk_size].copy() |
| chunk["text"] = chunk["text"].map(normalize_text) |
| chunk["label"] = chunk["label"].map(normalize_category) |
| chunk["source"] = chunk["source"].map(normalize_optional_text) |
|
|
| if "country" in chunk: |
| chunk["country"] = chunk["country"].map(normalize_country) |
| if "currency" in chunk: |
| chunk["currency"] = chunk["currency"].map(normalize_optional_text).str.upper() |
| if "language" in chunk: |
| chunk["language"] = chunk["language"].map(normalize_optional_text) |
| if "source_label" in chunk: |
| chunk["source_label"] = chunk["source_label"].map(normalize_category) |
| if "record_id" in chunk: |
| chunk["record_id"] = chunk["record_id"].map(normalize_optional_text) |
|
|
| |
| |
| |
| chunk = chunk[(chunk["text"] != "") & (chunk["label"] != "")] |
|
|
| if canonicalize: |
| mapped = [] |
| for source, label in zip(chunk["source"], chunk["label"], strict=True): |
| try: |
| mapped.append(canonicalize_category(label, source)) |
| except ValueError: |
| unknown_labels.add((source, label)) |
| mapped.append("") |
| chunk["label"] = mapped |
|
|
| chunk = chunk[(chunk["text"] != "") & (chunk["label"] != "")] |
| chunks.append(chunk) |
| progress_bar.set_postfix(rows=min(start + chunk_size, len(result))) |
|
|
| if unknown_labels and strict_taxonomy: |
| examples = ", ".join(f"{source}:{label}" for source, label in sorted(unknown_labels)[:25]) |
| raise ValueError( |
| "Training data contains source labels without an explicit canonical mapping: " |
| f"{examples}" |
| ) |
| if not chunks: |
| return result.iloc[0:0].reset_index(drop=True) |
|
|
| result = pd.concat(chunks, ignore_index=True) |
| dedupe_columns = ["text", "label", "source"] |
| for column in ("country", "currency", "language"): |
| if column in result.columns: |
| dedupe_columns.append(column) |
| result = result.drop_duplicates(subset=dedupe_columns, keep="first") |
| return result.reset_index(drop=True) |
|
|
