data_extraction.py

"""
NADRA Urdu -> Roman Urdu: Extract & Format for Fine-Tuning
============================================================

Extracts real Urdu/Roman-Urdu parallel pairs from HBL's NADRA transliteration
API logs and converts them into an instruction-tuned JSONL dataset ready for
fine-tuning tiny-aya-global.

Source of truth per record:
  - REQUEST_BODY -> original Urdu script (input)
  - data         -> the transliteration service's Roman Urdu output (target)

Fields used: name, fatherName, motherName, address1, address2, placeOfBirth

Usage:
    python nadra_data_prep.py --input "/path/to/*.json" --output-dir /path/to/output

Run with -h for all options.
"""

import argparse
import csv
import glob
import json
import os
import random
import re
import unicodedata
from collections import Counter

FIELDS = ["name", "fatherName", "motherName", "address1", "address2", "placeOfBirth"]

INSTRUCTIONS = {
    "name":         "Transliterate the following Urdu name into Roman Urdu.",
    "fatherName":   "Transliterate the following Urdu father's name into Roman Urdu.",
    "motherName":   "Transliterate the following Urdu mother's name into Roman Urdu.",
    "address1":     "Transliterate the following Urdu address into Roman Urdu.",
    "address2":     "Transliterate the following Urdu address into Roman Urdu.",
    "placeOfBirth": "Transliterate the following Urdu place name into Roman Urdu.",
}


# --------------------------------------------------------------------------
# Loading
# --------------------------------------------------------------------------
def _absorb(obj, records):
    """Given a parsed JSON value, pull out record dicts from known shapes."""
    if isinstance(obj, dict) and "response" in obj:
        records.extend(obj["response"])
    elif isinstance(obj, list):
        records.extend(obj)
    elif isinstance(obj, dict):
        records.append(obj)


def _parse_concatenated_json(raw):
    """
    Parses a string containing one or more back-to-back JSON values
    (e.g. multiple {"response": [...]} wrapper objects with no separator,
    or no enclosing list) using a streaming decoder.

    Skips standard whitespace AND invisible unicode separator artifacts
    (zero-width space, stray BOM) that sometimes appear between pages when
    multiple exports get concatenated into one file.

    Returns a list of parsed top-level JSON values.
    """
    INVISIBLE_SEPARATORS = ("\u200b", "\ufeff", "\u200c", "\u200d")
    decoder = json.JSONDecoder()
    values = []
    idx = 0
    n = len(raw)
    while idx < n:
        while idx < n and (raw[idx].isspace() or raw[idx] in INVISIBLE_SEPARATORS):
            idx += 1
        if idx >= n:
            break
        obj, end = decoder.raw_decode(raw, idx)
        values.append(obj)
        idx = end
    return values


def _context_snippet(raw, pos, window=120):
    start = max(0, pos - window)
    end = min(len(raw), pos + window)
    before = raw[start:pos]
    after = raw[pos:end]
    return f"...{before!r} <-- ERROR HERE --> {after!r}..."


def _preprocess_mongo_shell_syntax(raw):
    """
    Converts MongoDB *shell* export syntax into plain JSON.
    mongoexport/mongodump produce valid JSON (e.g. {"$oid": "..."}), but a
    raw copy-paste from the mongo shell or Compass often contains
    constructor-style tokens that are NOT valid JSON, e.g.:
        ObjectId("665f1...")   -> "665f1..."
        ISODate("2024-01-01T00:00:00Z") -> "2024-01-01T00:00:00Z"
        NumberLong(123)        -> 123
        NumberInt(123)         -> 123
        NumberDecimal("1.5")   -> "1.5"
        Timestamp(123, 1)      -> "Timestamp(123, 1)"  (rare; kept as string)
    This rewrites those into equivalent plain JSON values.
    """
    # ObjectId("...") / ISODate("...") / NumberDecimal("...") -> just the quoted string
    raw = re.sub(r'\b(?:ObjectId|ISODate|NumberDecimal|UUID)\(\s*"([^"]*)"\s*\)', r'"\1"', raw)
    # NumberLong(123) / NumberInt(123) -> bare number (handles quoted numbers too)
    raw = re.sub(r'\b(?:NumberLong|NumberInt)\(\s*"?(-?\d+)"?\s*\)', r'\1', raw)
    return raw


def load_records(filepaths):
    """
    Handles several shapes robustly:
      1) {"response": [ {...}, {...}, ... ]}   (paginated API wrapper)
      2) [ {...}, {...}, ... ]                 (plain list of records)
      3) JSON Lines: one record dict per line
      4) Multiple JSON values concatenated back-to-back in one file
    Also strips BOM and skips empty files, reporting problems per-file
    instead of crashing the whole run.
    """
    records = []
    for fp in filepaths:
        with open(fp, "r", encoding="utf-8-sig") as fh:  # utf-8-sig strips BOM if present
            raw = fh.read().strip()

        if not raw:
            print(f"  [skip] {fp} is empty")
            continue

        raw = _preprocess_mongo_shell_syntax(raw)

        whole_file_error = None
        try:
            obj = json.loads(raw)
            _absorb(obj, records)
            continue
        except json.JSONDecodeError as e:
            whole_file_error = e

        concat_error = None
        try:
            values = _parse_concatenated_json(raw)
            if values:
                for v in values:
                    _absorb(v, records)
                print(f"  [ok] {fp}: parsed as {len(values)} concatenated JSON value(s)")
                continue
        except json.JSONDecodeError as e:
            concat_error = e

        jsonl_error = None
        try:
            line_records = []
            for line in raw.splitlines():
                line = line.strip()
                if not line:
                    continue
                line_records.append(json.loads(line))
            for v in line_records:
                _absorb(v, records)
            print(f"  [ok] {fp}: parsed as JSON Lines ({len(line_records)} lines)")
            continue
        except json.JSONDecodeError as e:
            jsonl_error = e

        # All strategies failed -- report the WHOLE-FILE error, since that's the
        # most informative one for a file that's meant to be a single JSON document.
        print(f"  [FAILED] {fp}: could not parse as JSON, concatenated JSON, or JSON Lines.")
        print(f"           Whole-file parse error: {whole_file_error}")
        print(f"           Context around error position:")
        print(f"           {_context_snippet(raw, whole_file_error.pos)}")
        print(f"           (concatenated-JSON attempt error: {concat_error})")
        print(f"           (JSON-Lines attempt error: {jsonl_error})")

    return records


# --------------------------------------------------------------------------
# Extraction
# --------------------------------------------------------------------------
def extract_pairs(records, fields):
    """
    For each record, pairs record['REQUEST_BODY'][field] (Urdu)
    with record['data'][field] (Roman Urdu), for each field in `fields`.
    Skips records missing either block entirely.
    """
    pairs = []
    skipped_missing_block = 0

    for rec in records:
        data_block = rec.get("data")
        req_block = rec.get("REQUEST_BODY")
        if not isinstance(data_block, dict) or not isinstance(req_block, dict):
            skipped_missing_block += 1
            continue

        cnic = data_block.get("cnic", "")
        for field in fields:
            roman = data_block.get(field)
            urdu = req_block.get(field)
            if roman is None or urdu is None:
                continue
            pairs.append({"field": field, "urdu": urdu, "roman": roman, "cnic": cnic})

    print(f"Extracted {len(pairs)} raw pairs "
          f"({skipped_missing_block} records skipped for missing data/REQUEST_BODY blocks)")
    return pairs


# --------------------------------------------------------------------------
# Cleaning
# --------------------------------------------------------------------------
def is_masked_or_junk(text):
    if text is None:
        return True
    t = text.strip()
    if t == "":
        return True
    if "#" in t:  # redacted/placeholder PII
        return True
    if re.fullmatch(r"[\W_]+", t):  # only punctuation/symbols, no real content
        return True
    return False


def normalize_text(text):
    t = unicodedata.normalize("NFC", text)
    t = re.sub(r"\s+", " ", t).strip()
    return t


def clean_pairs(pairs):
    cleaned = []
    seen = set()
    dropped_masked = 0
    dropped_dupe = 0

    for p in pairs:
        urdu, roman = p["urdu"], p["roman"]

        if is_masked_or_junk(urdu) or is_masked_or_junk(roman):
            dropped_masked += 1
            continue

        urdu_n = normalize_text(urdu)
        roman_n = normalize_text(roman)

        key = (p["field"], urdu_n.lower(), roman_n.lower())
        if key in seen:
            dropped_dupe += 1
            continue
        seen.add(key)

        cleaned.append({"field": p["field"], "urdu": urdu_n, "roman": roman_n})

    print(f"Kept {len(cleaned)} pairs "
          f"(dropped {dropped_masked} masked/junk, {dropped_dupe} exact duplicates)")
    return cleaned


# --------------------------------------------------------------------------
# Formatting
# --------------------------------------------------------------------------
def to_instruction_example(pair):
    return {
        "instruction": INSTRUCTIONS[pair["field"]],
        "input": pair["urdu"],
        "output": pair["roman"],
        "field": pair["field"],  # kept for filtering/analysis; drop before training if not needed
    }


def train_val_split(examples, val_fraction, seed):
    random.seed(seed)
    by_field = {}
    for ex in examples:
        by_field.setdefault(ex["field"], []).append(ex)

    train_set, val_set = [], []
    for field, items in by_field.items():
        items = items[:]
        random.shuffle(items)
        n_val = max(1, int(len(items) * val_fraction)) if len(items) > 20 else 0
        val_set.extend(items[:n_val])
        train_set.extend(items[n_val:])

    random.shuffle(train_set)
    random.shuffle(val_set)
    return train_set, val_set


def write_jsonl(path, rows):
    with open(path, "w", encoding="utf-8") as fh:
        for row in rows:
            output_row = {k: v for k, v in row.items() if k != "field"}
            fh.write(json.dumps(output_row, ensure_ascii=False) + "\n")


def write_csv(path, pairs):
    with open(path, "w", newline="", encoding="utf-8") as fh:
        writer = csv.DictWriter(fh, fieldnames=["urdu", "roman"])
        writer.writeheader()
        for p in pairs:
            writer.writerow({"urdu": p["urdu"], "roman": p["roman"]})


# --------------------------------------------------------------------------
# Main
# --------------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser(description="Extract & format NADRA Urdu/Roman Urdu pairs for fine-tuning.")
    parser.add_argument("--input", required=True,
                         help='Path or glob pattern to input JSON file(s), e.g. "/data/nadra_page_*.json"')
    parser.add_argument("--output-dir", default="./output", help="Directory to write train.jsonl/val.jsonl/csv to")
    parser.add_argument("--val-fraction", type=float, default=0.05, help="Fraction held out for validation")
    parser.add_argument("--seed", type=int, default=42, help="Random seed for shuffling/splitting")
    parser.add_argument("--sample", type=int, default=3, help="Number of example pairs to print per field")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)
    train_path = os.path.join(args.output_dir, "train.jsonl")
    val_path = os.path.join(args.output_dir, "val.jsonl")
    csv_path = os.path.join(args.output_dir, "pairs_clean.csv")

    files = sorted(glob.glob(args.input))
    print(f"Found {len(files)} file(s) matching input pattern")
    for f in files:
        print(" -", f)
    if not files:
        print("No input files found. Check --input path/glob.")
        return

    records = load_records(files)
    print(f"Loaded {len(records)} raw records")

    raw_pairs = extract_pairs(records, FIELDS)
    print("Per-field raw counts:", Counter(p["field"] for p in raw_pairs))

    clean = clean_pairs(raw_pairs)
    print("Per-field clean counts:", Counter(p["field"] for p in clean))

    write_csv(csv_path, clean)
    print(f"Wrote clean pairs CSV -> {csv_path}")

    # Sanity print
    by_field = {}
    for p in clean:
        by_field.setdefault(p["field"], []).append(p)
    random.seed(args.seed)
    for field, items in by_field.items():
        print(f"\n--- {field} ({len(items)} pairs) ---")
        for ex in random.sample(items, min(args.sample, len(items))):
            print(f"  UR: {ex['urdu']}")
            print(f"  RO: {ex['roman']}")

    examples = [to_instruction_example(p) for p in clean]
    train_set, val_set = train_val_split(examples, args.val_fraction, args.seed)

    write_jsonl(train_path, train_set)
    write_jsonl(val_path, val_set)
    print(f"\nWrote {len(train_set)} rows -> {train_path}")
    print(f"Wrote {len(val_set)} rows -> {val_path}")


if __name__ == "__main__":
    main()