training/add_new_data.py

#!/usr/bin/env python
"""
Add newly extracted subject files to the MCQ dataset, safely.

Why a script instead of `cat a b > c`: on Windows PowerShell the `>` redirect re-encodes the text and
corrupts Sinhala. This reads and writes UTF-8 explicitly and validates every line.

It figures out each input file's format on its own:
  * raw extraction format  (original_question / answer_1..5 / correct_answer)  -> runs mcq_dataset
    convert on it (same cleaning, junk-stripping and filters as the original pipeline)
  * MCQ schema             (question / choices / answer)                        -> passed through
It then merges everything into one converted file, ready for clean -> dedupe -> split -> build-sft.

    python add_new_data.py --base data/converted.jsonl \
        --new data/art.jsonl data/buddhim.jsonl data/dancing.jsonl data/sinhala.jsonl \
        --out data/all_converted.jsonl

Prints a per-file report (rows kept, subjects, exam types/grades) so you can see what you actually
added before spending GPU time. Put mcq_dataset.py in the same folder.
"""
import argparse
import collections
import importlib.util
import json
import os
import sys
from argparse import Namespace

HERE = os.path.dirname(os.path.abspath(__file__))


def load_mcq_module():
    path = os.path.join(HERE, "mcq_dataset.py")
    if not os.path.exists(path):
        sys.exit(f"mcq_dataset.py not found next to this script ({HERE})")
    spec = importlib.util.spec_from_file_location("mcq_dataset", path)
    m = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(m)
    return m


def read_jsonl(path):
    rows, bad = [], 0
    with open(path, encoding="utf-8-sig") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError:
                bad += 1
    if bad:
        print(f"  [warn] {bad} unparseable line(s) skipped in {os.path.basename(path)}")
    return rows


def detect_format(rows):
    if not rows:
        return "empty"
    r = rows[0]
    if "choices" in r and "answer" in r and "question" in r:
        return "mcq"
    if "original_question" in r or "answer_1" in r or "correct_answer" in r:
        return "raw"
    return "unknown"


def describe(rows, label):
    subj = collections.Counter(r.get("subject") or "?" for r in rows)
    lvl = collections.Counter()
    for r in rows:
        md = r.get("metadata") or {}
        lvl[str(md.get("exam_type") or r.get("level") or "?") + "/" +
            str(md.get("grade") or r.get("grade") or "?")] += 1
    nch = collections.Counter(len(r.get("choices") or []) for r in rows)
    print(f"  {label}: {len(rows)} rows")
    print(f"    subjects : {dict(subj.most_common(8))}")
    print(f"    level/grade: {dict(lvl.most_common(8))}")
    print(f"    n_choices: {dict(sorted(nch.items()))}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base", default="data/converted.jsonl",
                    help="existing converted file to append to ('' = start fresh)")
    ap.add_argument("--new", nargs="+", required=True, help="newly extracted subject files")
    ap.add_argument("--out", required=True)
    ap.add_argument("--tmp_dir", default="", help="where to write per-file converted output "
                                                  "(default: alongside each input)")
    ap.add_argument("--keep_latin", action="store_true")
    ap.add_argument("--latin_allow", nargs="*", default=[])
    ap.add_argument("--min_confidence", choices=["low", "medium", "high"], default="low")
    args = ap.parse_args()
    for stream in (sys.stdout, sys.stderr):
        try:
            stream.reconfigure(encoding="utf-8")
        except Exception:                                                # noqa: BLE001
            pass

    mcq = load_mcq_module()
    merged, seen_ids = [], set()

    if args.base:
        if not os.path.exists(args.base):
            sys.exit(f"--base {args.base} not found (pass --base '' to start fresh)")
        base_rows = read_jsonl(args.base)
        if detect_format(base_rows) != "mcq":
            sys.exit(f"--base {args.base} is not in MCQ schema; pass the CONVERTED file")
        print(f"[base] {args.base}")
        describe(base_rows, "existing")
        merged.extend(base_rows)

    for path in args.new:
        if not os.path.exists(path):
            print(f"[skip] {path} not found")
            continue
        rows = read_jsonl(path)
        fmt = detect_format(rows)
        print(f"\n[new] {os.path.basename(path)}  format={fmt}")
        if fmt == "raw":
            out_dir = args.tmp_dir or os.path.dirname(path) or "."
            os.makedirs(out_dir, exist_ok=True)
            stem = os.path.splitext(os.path.basename(path))[0]
            conv = os.path.join(out_dir, f"{stem}_conv.jsonl")
            rej = os.path.join(out_dir, f"{stem}_latin_rejects.jsonl")
            mcq.cmd_convert(Namespace(inp=path, out=conv, min_confidence=args.min_confidence,
                                      skip_needs_review=False, keep_shared_stem=False,
                                      keep_latin=args.keep_latin, latin_allow=args.latin_allow,
                                      rejects=rej))
            rows = read_jsonl(conv)
        elif fmt == "mcq":
            print("  already in MCQ schema: passed through without conversion")
        else:
            print(f"  [skip] unrecognised format; first keys: {sorted(rows[0])[:8] if rows else '-'}")
            continue
        describe(rows, "added")
        merged.extend(rows)

    # ids must be unique, or later stages (and any per-question analysis) get confused
    out_rows, dup = [], 0
    for i, r in enumerate(merged):
        rid = r.get("id")
        if rid is None or rid in seen_ids:
            dup += 1
            rid = f"{rid or 'row'}__{i}"
        seen_ids.add(rid)
        r["id"] = rid
        out_rows.append(r)
    if dup:
        print(f"\n[ids] {dup} missing/duplicate id(s) were made unique")

    os.makedirs(os.path.dirname(os.path.abspath(args.out)) or ".", exist_ok=True)
    with open(args.out, "w", encoding="utf-8") as f:
        for r in out_rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")

    print(f"\n[out] {len(out_rows)} rows -> {args.out}")
    describe(out_rows, "combined")
    print("\nNext (run from the same folder):")
    print(f"  python mcq_dataset.py validate {args.out}")
    print(f"  python mcq_dataset.py clean {args.out} data/all_cleaned.jsonl --rejects data/all_clean_rejects.jsonl")
    print("  python mcq_dataset.py dedupe data/all_cleaned.jsonl data/all_deduped.jsonl "
          "--blocklist test_easy_block.jsonl --conflicts data/all_conflicts.jsonl")
    print("  python mcq_dataset.py stats data/all_deduped.jsonl")
    print("  python mcq_dataset.py split data/all_deduped.jsonl --train data/train2.jsonl "
          "--val data/val2.jsonl --val_frac 0.1")
    print("  python mcq_dataset.py build-sft data/train2.jsonl data/sft_train2.jsonl --shuffle_copies 2")
    print("  python mcq_dataset.py build-sft data/val2.jsonl data/sft_val2.jsonl "
          "--shuffle_copies 0 --strip_zwj_prob 0 --templates si_instruction")


if __name__ == "__main__":
    main()