Create data_extraction.py

data_extraction.py

@@ -0,0 +1,262 @@
+"""
+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 load_records(filepaths):
+    """
+    Handles three shapes:
+      1) {"response": [ {...}, {...}, ... ]}   (paginated API wrapper)
+      2) [ {...}, {...}, ... ]                 (plain list of records)
+      3) JSON Lines: one record dict per line
+    Returns a flat list of record dicts.
+    """
+    records = []
+    for fp in filepaths:
+        with open(fp, "r", encoding="utf-8") as fh:
+            raw = fh.read().strip()
+        try:
+            obj = json.loads(raw)
+            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)
+        except json.JSONDecodeError:
+            # fall back to JSON Lines
+            for line in raw.splitlines():
+                line = line.strip()
+                if not line:
+                    continue
+                records.append(json.loads(line))
+    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:
+            fh.write(json.dumps(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=["field", "urdu", "roman"])
+        writer.writeheader()
+        writer.writerows(pairs)
+
+
+# --------------------------------------------------------------------------
+# 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()