combine_data.py

#!/usr/bin/env python3
"""
Combine training data from multiple sources and create train/valid split.
"""

import json
import random
from pathlib import Path

# Input files
TWITTER_DATA = Path("train.jsonl")
BLUESKY_DATA = Path("bluesky_posts.jsonl")

# Output directory
OUTPUT_DIR = Path("combined_data")

# Validation split ratio
VALID_RATIO = 0.05


def load_jsonl(path: Path) -> list[dict]:
    """Load a JSONL file."""
    data = []
    with open(path) as f:
        for line in f:
            if line.strip():
                data.append(json.loads(line))
    return data


def main():
    OUTPUT_DIR.mkdir(exist_ok=True)

    # Load all data
    print(f"Loading {TWITTER_DATA}...")
    twitter_data = load_jsonl(TWITTER_DATA)
    print(f"  Loaded {len(twitter_data)} Twitter examples")

    print(f"Loading {BLUESKY_DATA}...")
    bluesky_data = load_jsonl(BLUESKY_DATA)
    print(f"  Loaded {len(bluesky_data)} Bluesky examples")

    # Combine
    all_data = twitter_data + bluesky_data
    print(f"\nTotal combined: {len(all_data)} examples")

    # Shuffle
    random.seed(42)
    random.shuffle(all_data)

    # Split
    valid_size = int(len(all_data) * VALID_RATIO)
    valid_data = all_data[:valid_size]
    train_data = all_data[valid_size:]

    print(f"Train: {len(train_data)} examples")
    print(f"Valid: {len(valid_data)} examples")

    # Write
    train_path = OUTPUT_DIR / "train.jsonl"
    valid_path = OUTPUT_DIR / "valid.jsonl"

    with open(train_path, "w") as f:
        for item in train_data:
            f.write(json.dumps(item) + "\n")

    with open(valid_path, "w") as f:
        for item in valid_data:
            f.write(json.dumps(item) + "\n")

    print(f"\nWritten to {OUTPUT_DIR}/")
    print(f"  {train_path}: {train_path.stat().st_size / 1024 / 1024:.1f} MB")
    print(f"  {valid_path}: {valid_path.stat().st_size / 1024 / 1024:.1f} MB")


if __name__ == "__main__":
    main()