sample_semantics4coding.py

import json
import random
from pathlib import Path

ROOT = Path("/data01/RAG4Coding_datasets/Semantics4Coding")
OUTPUT = ROOT / "sample_6000.jsonl"
SAMPLE_PER_SOURCE = 1000

random.seed(42)

def normalize(example):
    # Case 1: already prompt + output
    if "prompt" in example and "output" in example:
        return {
            "prompt": example["prompt"],
            "output": example["output"],
        }
    # Case 2: input + output → prompt + output
    if "input" in example and "output" in example:
        return {
            "prompt": example["input"],
            "output": example["output"],
        }
    # Case 3: TCU format → prompt + test_case
    if "prompt" in example and "test_case" in example:
        return {
            "prompt": example["prompt"],
            "output": example["test_case"],
        }
    return None


def read_and_normalize(path):
    out = []
    with open(path, "r") as f:
        for line in f:
            ex = normalize(json.loads(line))
            if ex is not None:
                out.append(ex)
    return out

all_samples = []

counts = {}

for item in ROOT.iterdir():
    if item.name.startswith("sample_"):
        continue

    name = item.name

    if item.is_file() and item.suffix == ".jsonl":
        data = read_and_normalize(item)
        take = min(SAMPLE_PER_SOURCE, len(data))
        all_samples.extend(random.sample(data, take))
        counts[name] = take

    elif item.is_dir():
        merged = []
        for jsonl_file in item.glob("*.jsonl"):
            merged.extend(read_and_normalize(jsonl_file))

        if merged:
            take = min(SAMPLE_PER_SOURCE, len(merged))
            all_samples.extend(random.sample(merged, take))
            counts[name] = take

with open(OUTPUT, "w") as f:
    for ex in all_samples:
        f.write(json.dumps(ex) + "\n")

print("Samples per source:")
for k, v in counts.items():
    print(f"{k}: {v}")

print(f"\nTotal samples written: {len(all_samples)}")
print(f"Output: {OUTPUT}")