File size: 3,093 Bytes
877049d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | import argparse
from pathlib import Path
import pandas as pd
from utils import clean_title, ensure_dir, normalize_label, save_json
def load_input(path: str | Path) -> pd.DataFrame:
path = Path(path)
if path.suffix.lower() == ".csv":
return pd.read_csv(path)
if path.suffix.lower() == ".jsonl":
return pd.read_json(path, lines=True)
raise ValueError(f"Unsupported input format: {path.suffix}")
def preprocess_frame(dataframe: pd.DataFrame, min_title_length: int) -> tuple[pd.DataFrame, dict]:
stats = {"before_rows": int(len(dataframe))}
if "title" not in dataframe.columns or "label" not in dataframe.columns:
raise ValueError("Input data must contain `title` and `label` columns.")
processed = dataframe.copy()
processed["title"] = processed["title"].map(clean_title)
processed = processed[processed["title"].str.len() >= min_title_length].copy()
stats["after_title_filter"] = int(len(processed))
processed["label_id"] = processed["label"].map(normalize_label)
processed["label"] = processed["label_id"].map({0: "non-clickbait", 1: "clickbait"})
processed = processed.drop_duplicates(subset=["title"]).reset_index(drop=True)
stats["after_deduplicate"] = int(len(processed))
processed["node_id"] = processed.index.astype(int)
preferred_columns = ["node_id", "id", "title", "label", "label_id", "source", "category", "publish_datetime", "url"]
existing_columns = [column for column in preferred_columns if column in processed.columns]
remaining_columns = [column for column in processed.columns if column not in existing_columns]
processed = processed[existing_columns + remaining_columns]
stats["label_distribution"] = processed["label"].value_counts().to_dict()
return processed, stats
def main() -> None:
parser = argparse.ArgumentParser(description="Preprocess ViClickbait dataset.")
parser.add_argument("--input", required=True, help="Path to raw CSV or JSONL.")
parser.add_argument("--output", required=True, help="Path to write processed CSV.")
parser.add_argument("--label_map", default=None, help="Optional label_map.json path.")
parser.add_argument("--stats_path", default=None, help="Optional preprocess_stats.json path.")
parser.add_argument("--min_title_length", type=int, default=4)
args = parser.parse_args()
frame = load_input(args.input)
processed, stats = preprocess_frame(frame, min_title_length=args.min_title_length)
output_path = Path(args.output)
ensure_dir(output_path.parent)
processed.to_csv(output_path, index=False)
label_map_path = Path(args.label_map) if args.label_map else output_path.parent / "label_map.json"
save_json({"label_to_id": {"non-clickbait": 0, "clickbait": 1}}, label_map_path)
stats_path = Path(args.stats_path) if args.stats_path else output_path.parent / "preprocess_stats.json"
save_json(stats, stats_path)
print(f"Saved processed data to {output_path}")
print(f"Rows after preprocessing: {len(processed)}")
if __name__ == "__main__":
main()
|