scripts/validate_jsonl.py

#!/usr/bin/env python3
"""
validate_jsonl.py

Validates all JSONL rows in the dataset files.
Invalid rows are moved to datasets/mythos_coder_rejected.jsonl.

Usage:
    python validate_jsonl.py [--train-only] [--valid-only]

Returns exit code 0 if all valid, 1 if any invalid.
"""

import argparse
import json
import sys
from pathlib import Path


def load_schema():
    """Load the JSON schema for validation."""
    schema_path = Path(__file__).parent.parent / "schemas" / "training_example.schema.json"
    with open(schema_path, "r", encoding="utf-8") as f:
        return json.load(f)


def validate_example(example, schema):
    """Validate a single example against the schema. Returns list of errors."""
    required = schema.get("required", [])
    errors = []

    for field in required:
        if field not in example:
            errors.append(f"Missing required field: {field}")

    if "task_type" in example:
        allowed_types = schema["properties"]["task_type"]["enum"]
        if example["task_type"] not in allowed_types:
            errors.append(f"Invalid task_type: {example['task_type']}")

    if "difficulty" in example:
        allowed_difficulties = schema["properties"]["difficulty"]["enum"]
        if example["difficulty"] not in allowed_difficulties:
            errors.append(f"Invalid difficulty: {example['difficulty']}")

    if "quality_score" in example:
        score = example["quality_score"]
        if not isinstance(score, int) or score < 1 or score > 10:
            errors.append(f"Invalid quality_score: {score} (must be 1-10)")

    if "investigation_steps" in example:
        if not isinstance(example["investigation_steps"], list):
            errors.append("investigation_steps must be an array")

    return errors


def validate_file(file_path, schema, rejected_path=None):
    """
    Validate all lines in a JSONL file.
    Returns (valid_count, invalid_count, invalid_lines).
    If rejected_path is provided, invalid lines are written there.
    """
    if not file_path.exists():
        return 0, 0, []

    valid_count = 0
    invalid_count = 0
    invalid_lines = []
    rejected_examples = []

    with open(file_path, "r", encoding="utf-8") as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue

            try:
                example = json.loads(line)
                errors = validate_example(example, schema)

                if errors:
                    invalid_count += 1
                    invalid_lines.append((line_num, errors, line))
                    rejected_examples.append(example)
                else:
                    valid_count += 1
            except json.JSONDecodeError as e:
                invalid_count += 1
                invalid_lines.append((line_num, [f"JSON parse error: {e}"], line))
                rejected_examples.append({"_raw": line, "_error": str(e)})

    # Write rejected examples to rejected file
    if rejected_path and rejected_examples:
        rejected_path.parent.mkdir(parents=True, exist_ok=True)
        with open(rejected_path, "a", encoding="utf-8") as f:
            for example in rejected_examples:
                f.write(json.dumps(example, ensure_ascii=False) + "\n")

    return valid_count, invalid_count, invalid_lines


def main():
    parser = argparse.ArgumentParser(description="Validate JSONL dataset files")
    parser.add_argument("--train-only", action="store_true", help="Only validate training file")
    parser.add_argument("--valid-only", action="store_true", help="Only validate validation file")
    parser.add_argument("--verbose", "-v", action="store_true", help="Show all validation errors")
    args = parser.parse_args()

    project_root = Path(__file__).parent.parent
    schema = load_schema()

    train_path = project_root / "datasets" / "mythos_coder_train.jsonl"
    valid_path = project_root / "datasets" / "mythos_coder_valid.jsonl"
    rejected_path = project_root / "datasets" / "mythos_coder_rejected.jsonl"

    all_valid = True
    total_valid = 0
    total_invalid = 0

    # Validate training file
    if not args.valid_only:
        print(f"Validating {train_path}...")
        valid, invalid, errors = validate_file(train_path, schema, rejected_path if invalid > 0 else None)
        total_valid += valid
        total_invalid += invalid
        print(f"  Valid: {valid}, Invalid: {invalid}")

        if args.verbose and errors:
            for line_num, errs, _ in errors:
                print(f"    Line {line_num}: {errs}")

        if invalid > 0:
            print(f"  Invalid examples moved to {rejected_path}")
            all_valid = False

    # Validate validation file
    if not args.train_only:
        print(f"\nValidating {valid_path}...")
        valid, invalid, errors = validate_file(valid_path, schema, rejected_path if invalid > 0 else None)
        total_valid += valid
        total_invalid += invalid
        print(f"  Valid: {valid}, Invalid: {invalid}")

        if args.verbose and errors:
            for line_num, errs, _ in errors:
                print(f"    Line {line_num}: {errs}")

        if invalid > 0:
            print(f"  Invalid examples moved to {rejected_path}")
            all_valid = False

    print(f"\n{'='*50}")
    print(f"Total: {total_valid} valid, {total_invalid} invalid")

    sys.exit(0 if all_valid else 1)


if __name__ == "__main__":
    main()