Upload format_dataset.py

format_dataset.py

@@ -0,0 +1,133 @@
+#!/usr/bin/env python3
+"""
+Format BothBosu scam-dialogue (and optionally Scammer-Conversation)
+into standardized chat-template JSONL for SFT.
+
+REQUIREMENTS:
+  pip install datasets transformers
+
+USAGE:
+  python format_dataset.py --out_dir ./formatted_scam_data
+
+OUTPUT:
+  formatted_scam_data/
+    train.jsonl   (chat-format messages per line)
+    test.jsonl
+    README.md     (dataset card fragment)
+
+Each JSONL line:
+  {"messages": [
+    {"role": "system",    "content": "You are a phone scam detection expert."},
+    {"role": "user",      "content": "Read this transcript...\n\n{transcript}"},
+    {"role": "assistant", "content": "SCAM"}
+  ]}
+"""
+
+import argparse
+import json
+from pathlib import Path
+
+from datasets import load_dataset, concatenate_datasets
+
+
+PROMPT_TEMPLATE = (
+    "Read this phone call transcript and classify it:\n\n"
+    "{transcript}\n\n"
+    "Answer with exactly ONE word: SCAM or LEGITIMATE."
+)
+SYSTEM = "You are a phone scam detection expert."
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--primary", default="BothBosu/scam-dialogue")
+    p.add_argument("--secondary", default="BothBosu/Scammer-Conversation",
+                   help="Optional extra dataset to merge into train")
+    p.add_argument("--out_dir", default="./formatted_scam_data")
+    return p.parse_args()
+
+
+def row_to_chat(row):
+    """Convert a raw dataset row → ChatML dict."""
+    answer = "SCAM" if row["label"] == 1 else "LEGITIMATE"
+    # Handle different column names across datasets
+    transcript = row.get("dialogue") or row.get("conversation")
+    return {
+        "messages": [
+            {"role": "system",    "content": SYSTEM},
+            {"role": "user",      "content": PROMPT_TEMPLATE.format(transcript=transcript)},
+            {"role": "assistant", "content": answer},
+        ]
+    }
+
+
+def save_jsonl(rows, path: Path):
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with open(path, "w", encoding="utf-8") as f:
+        for r in rows:
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")
+    print(f"Saved {len(rows)} rows → {path}")
+
+
+def main():
+    args = parse_args()
+    out_dir = Path(args.out_dir)
+
+    # Load primary
+    print(f"Loading primary dataset: {args.primary}")
+    ds_train = load_dataset(args.primary, split="train")
+    ds_test  = load_dataset(args.primary, split="test")
+
+    # Optional secondary merge
+    if args.secondary:
+        try:
+            ds_extra = load_dataset(args.secondary, split="train")
+            n_before = len(ds_train)
+            ds_train = concatenate_datasets([ds_train, ds_extra])
+            print(f"Merged {args.secondary}: {n_before} → {len(ds_train)} train rows")
+        except Exception as e:
+            print(f"Skipped secondary dataset: {e}")
+
+    # Convert
+    train_rows = [row_to_chat(r) for r in ds_train]
+    test_rows  = [row_to_chat(r) for r in ds_test]
+
+    # Save
+    save_jsonl(train_rows, out_dir / "train.jsonl")
+    save_jsonl(test_rows,  out_dir / "test.jsonl")
+
+    # Stats
+    n_scam_train = sum(1 for r in train_rows if r["messages"][2]["content"] == "SCAM")
+    n_scam_test  = sum(1 for r in test_rows  if r["messages"][2]["content"] == "SCAM")
+
+    stats = {
+        "train": {"total": len(train_rows), "scam": n_scam_train, "legit": len(train_rows) - n_scam_train},
+        "test":  {"total": len(test_rows),  "scam": n_scam_test,  "legit": len(test_rows)  - n_scam_test},
+    }
+
+    (out_dir / "stats.json").write_text(json.dumps(stats, indent=2))
+    print(f"\nStats:\n{json.dumps(stats, indent=2)}")
+
+    # README fragment
+    readme = f"""# Formatted Scam-Call Dataset (ChatML)
+
+Generated by `format_dataset.py`.
+
+## Sources
+- Primary: {args.primary}
+- Secondary: {args.secondary or "None"}
+
+## Statistics
+```json
+{json.dumps(stats, indent=2)}
+```
+
+## Schema
+Each `.jsonl` line is a ChatML message list compatible with TRL / Unsloth SFTTrainer.
+"""
+    (out_dir / "README.md").write_text(readme)
+    print(f"\nDone. Output directory: {out_dir.absolute()}")
+
+
+if __name__ == "__main__":
+    main()