scripts/score_dataset.py

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

Analyzes dataset balance and quality metrics.

Usage:
    python score_dataset.py
    python score_dataset.py --file datasets/mythos_coder_valid.jsonl
    python score_dataset.py --min-quality 7
"""

import argparse
import json
from collections import defaultdict
from pathlib import Path


def load_examples(file_path):
    """Load all examples from a JSONL file."""
    examples = []
    if not file_path.exists():
        return examples

    with open(file_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    examples.append(json.loads(line))
                except json.JSONDecodeError:
                    continue
    return examples


def calculate_balance(counts, total):
    """Calculate balance score (1.0 = perfectly balanced)."""
    if total == 0:
        return 0.0

    # Using entropy-based balance metric
    import math
    entropy = 0.0
    for count in counts.values():
        if count > 0:
            p = count / total
            entropy -= p * math.log(p)

    # Normalize by max possible entropy
    num_categories = len(counts)
    if num_categories <= 1:
        return 1.0

    max_entropy = math.log(num_categories)
    return entropy / max_entropy if max_entropy > 0 else 0.0


def analyze_dataset(examples):
    """Analyze dataset composition and quality."""
    if not examples:
        return {}

    stats = {
        "total": len(examples),
        "by_task_type": defaultdict(int),
        "by_difficulty": defaultdict(int),
        "by_language": defaultdict(int),
        "by_framework": defaultdict(int),
        "quality_scores": [],
        "avg_quality": 0.0,
        "quality_distribution": defaultdict(int),
        "high_quality_count": 0,  # score >= 8
    }

    for ex in examples:
        stats["by_task_type"][ex.get("task_type", "unknown")] += 1
        stats["by_difficulty"][ex.get("difficulty", "unknown")] += 1
        stats["by_language"][ex.get("language", "unknown")] += 1
        stats["by_framework"][ex.get("framework", "unknown")] += 1

        score = ex.get("quality_score", 0)
        stats["quality_scores"].append(score)
        stats["quality_distribution"][score] += 1

        if score >= 8:
            stats["high_quality_count"] += 1

    stats["avg_quality"] = sum(stats["quality_scores"]) / len(stats["quality_scores"])

    # Calculate balance scores
    stats["task_type_balance"] = calculate_balance(stats["by_task_type"], stats["total"])
    stats["difficulty_balance"] = calculate_balance(stats["by_difficulty"], stats["total"])

    return stats


def print_stats(stats, label):
    """Print statistics in a readable format."""
    if not stats:
        print(f"\n{label}: No data")
        return

    print(f"\n{'='*60}")
    print(f"{label}")
    print(f"{'='*60}")
    print(f"Total examples: {stats['total']}")
    print(f"Average quality score: {stats['avg_quality']:.2f}")
    print(f"High quality (>=8): {stats['high_quality_count']} ({100*stats['high_quality_count']/stats['total']:.1f}%)")

    print(f"\nTask Type Distribution (balance: {stats['task_type_balance']:.2f}):")
    for task_type, count in sorted(stats["by_task_type"].items()):
        pct = 100 * count / stats["total"]
        bar = "█" * int(pct / 2)
        print(f"  {task_type:20s}: {count:4d} ({pct:5.1f}%) {bar}")

    print(f"\nDifficulty Distribution (balance: {stats['difficulty_balance']:.2f}):")
    order = ["beginner", "intermediate", "advanced", "expert", "unknown"]
    for difficulty in order:
        if difficulty in stats["by_difficulty"]:
            count = stats["by_difficulty"][difficulty]
            pct = 100 * count / stats["total"]
            bar = "█" * int(pct / 2)
            print(f"  {difficulty:20s}: {count:4d} ({pct:5.1f}%) {bar}")

    print(f"\nTop Languages:")
    for lang, count in sorted(stats["by_language"].items(), key=lambda x: -x[1])[:5]:
        pct = 100 * count / stats["total"]
        print(f"  {lang:20s}: {count:4d} ({pct:.1f}%)")

    print(f"\nTop Frameworks:")
    for fw, count in sorted(stats["by_framework"].items(), key=lambda x: -x[1])[:5]:
        pct = 100 * count / stats["total"]
        print(f"  {fw:20s}: {count:4d} ({pct:.1f}%)")

    print(f"\nQuality Score Distribution:")
    for score in sorted(stats["quality_distribution"].keys()):
        count = stats["quality_distribution"][score]
        pct = 100 * count / stats["total"]
        bar = "█" * count
        print(f"  Score {score}: {count:4d} ({pct:5.1f}%) {bar}")


def main():
    parser = argparse.ArgumentParser(description="Score dataset balance and quality")
    parser.add_argument("--file", "-f", help="Specific file to analyze (default: all)")
    parser.add_argument("--min-quality", "-q", type=int, help="Filter by minimum quality score")
    parser.add_argument("--compare", "-c", action="store_true", help="Compare train vs valid")
    args = parser.parse_args()

    project_root = Path(__file__).parent.parent

    if args.file:
        file_paths = {"Dataset": project_root / args.file}
    elif args.compare:
        file_paths = {
            "Training": project_root / "datasets" / "mythos_coder_train.jsonl",
            "Validation": project_root / "datasets" / "mythos_coder_valid.jsonl",
        }
    else:
        file_paths = {
            "Training": project_root / "datasets" / "mythos_coder_train.jsonl",
            "Validation": project_root / "datasets" / "mythos_coder_valid.jsonl",
            "Rejected": project_root / "datasets" / "mythos_coder_rejected.jsonl",
        }

    for label, path in file_paths.items():
        examples = load_examples(path)

        if args.min_quality is not None:
            examples = [ex for ex in examples if ex.get("quality_score", 0) >= args.min_quality]

        stats = analyze_dataset(examples)
        print_stats(stats, f"{label}: {path.name if hasattr(path, 'name') else args.file}")

    # Overall summary if comparing
    if args.compare and len(file_paths) == 2:
        train_examples = load_examples(file_paths["Training"])
        valid_examples = load_examples(file_paths["Validation"])

        train_stats = analyze_dataset(train_examples)
        valid_stats = analyze_dataset(valid_examples)

        if train_stats and valid_stats:
            print(f"\n{'='*60}")
            print("SPLIT COMPARISON")
            print(f"{'='*60}")

            train_set = set(ex.get("id") for ex in train_examples)
            valid_set = set(ex.get("id") for ex in valid_examples)
            overlap = train_set & valid_set

            if overlap:
                print(f"WARNING: {len(overlap)} IDs exist in both sets!")
            else:
                print("OK: No ID overlap between train and validation sets")

            # Compare distributions
            print(f"\nDistribution similarity:")
            for task_type in set(train_stats["by_task_type"].keys()) | set(valid_stats["by_task_type"].keys()):
                train_pct = 100 * train_stats["by_task_type"].get(task_type, 0) / train_stats["total"]
                valid_pct = 100 * valid_stats["by_task_type"].get(task_type, 0) / valid_stats["total"]
                diff = abs(train_pct - valid_pct)
                status = "OK" if diff < 5 else "WARN" if diff < 10 else "MISMATCH"
                print(f"  {task_type:15s}: Train {train_pct:5.1f}% | Valid {valid_pct:5.1f}% | {status}")


if __name__ == "__main__":
    main()