File size: 9,180 Bytes
875e4af
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
"""
Phase 3, Steps 1-3: inventory the real uploaded Parquet shard
(pipecat-ai/smart-turn-data-v3.2-train, shard 0 of 10, uploaded by the user
to /mnt/user-data/uploads/train-00000-of-00010.parquet) and materialize a
reproducible stratified sample of real audio for pipeline validation.

Uses the hand-written pure-Python Parquet reader in
src/turn_detector/parquet_reader.py, built specifically because this
sandbox has no network access to install pyarrow/fastparquet/datasets (see
docs/PHASE3_REAL_AUDIO_VALIDATION.md for the full story, including the
real bugs caught and fixed while building that reader against this file).
"""

from __future__ import annotations

import json
import subprocess
import sys
from collections import Counter
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from turn_detector.parquet_reader import (
    read_footer_metadata, read_column_values, parse_flac_streaminfo,
)
from turn_detector.data import stratified_reservoir_sample, duration_bucket

SHARD_PATH = "/mnt/user-data/uploads/train-00000-of-00010.parquet"
OUT_DIR = Path("data/raw/phase3_sample")
SCALAR_COLUMNS = ["id", "language", "endpoint_bool", "midfiller", "endfiller", "synthetic", "dataset"]


def decode_full_shard_metadata(shard_path: str) -> list[dict]:
    """Decode every scalar metadata column + audio duration (via FLAC
    STREAMINFO header only, no full audio decode) for EVERY row in the
    shard. This is cheap (~5s total, verified) since it doesn't touch the
    compressed audio payload beyond decompressing it to read a 34-byte
    header. Returns one dict per row, in file order.
    """
    pf = read_footer_metadata(shard_path)
    print(f"Shard: {pf.num_rows} total rows, {len(pf.row_groups)} row groups", file=sys.stderr)

    all_records = []
    corrupted = []

    for rg_idx, rg in enumerate(pf.row_groups):
        col_values = {}
        for colname in SCALAR_COLUMNS:
            col = rg.columns[colname]
            mdl = pf.max_def_levels[colname]
            col_values[colname] = read_column_values(shard_path, col, max_rows=rg.num_rows, max_def_level=mdl)

        audio_col = rg.columns["audio.bytes"]
        audio_mdl = pf.max_def_levels["audio.bytes"]
        audio_bytes_list = read_column_values(shard_path, audio_col, max_rows=rg.num_rows, max_def_level=audio_mdl)

        for i in range(rg.num_rows):
            rec = {c: col_values[c][i] for c in SCALAR_COLUMNS}
            for k in ("id", "language", "dataset"):
                if rec[k] is not None:
                    rec[k] = rec[k].decode("utf-8", errors="replace")
            audio_bytes = audio_bytes_list[i]
            if audio_bytes is None:
                corrupted.append({"row_group": rg_idx, "index_in_group": i, "id": rec["id"], "reason": "audio.bytes is null"})
                rec["duration_sec"] = None
                rec["sample_rate"] = None
                rec["channels"] = None
                rec["_corrupted"] = True
            else:
                try:
                    info = parse_flac_streaminfo(audio_bytes)
                    rec["duration_sec"] = info["duration_sec"]
                    rec["sample_rate"] = info["sample_rate"]
                    rec["channels"] = info["channels"]
                    rec["_corrupted"] = False
                except Exception as e:
                    corrupted.append({"row_group": rg_idx, "index_in_group": i, "id": rec["id"], "reason": str(e)})
                    rec["duration_sec"] = None
                    rec["sample_rate"] = None
                    rec["channels"] = None
                    rec["_corrupted"] = True
            rec["_row_group"] = rg_idx
            rec["_index_in_group"] = i
            all_records.append(rec)

        if (rg_idx + 1) % 8 == 0:
            print(f"  processed {rg_idx + 1}/{len(pf.row_groups)} row groups...", file=sys.stderr)

    print(f"Decoded {len(all_records)} rows. Corrupted/unreadable: {len(corrupted)}", file=sys.stderr)
    return all_records, corrupted, pf


def summarize(records: list[dict]) -> dict:
    n = len(records)
    def dist(key):
        c = Counter(str(r.get(key)) for r in records)
        return dict(sorted(c.items(), key=lambda kv: -kv[1]))

    durations = [r["duration_sec"] for r in records if r.get("duration_sec") is not None]
    sample_rates = Counter(r["sample_rate"] for r in records if r.get("sample_rate") is not None)
    channels = Counter(r["channels"] for r in records if r.get("channels") is not None)

    return {
        "n_rows": n,
        "endpoint_bool_distribution": dist("endpoint_bool"),
        "language_distribution": dist("language"),
        "dataset_source_distribution": dist("dataset"),
        "synthetic_distribution": dist("synthetic"),
        "midfiller_distribution_including_null": dist("midfiller"),
        "endfiller_distribution_including_null": dist("endfiller"),
        "sample_rates_found": dict(sample_rates),
        "channel_counts_found": dict(channels),
        "n_corrupted": sum(1 for r in records if r.get("_corrupted")),
        "duration_stats_sec": {
            "n": len(durations),
            "min": min(durations) if durations else None,
            "max": max(durations) if durations else None,
            "mean": sum(durations) / len(durations) if durations else None,
            "median": sorted(durations)[len(durations) // 2] if durations else None,
        },
    }


def main():
    print("=" * 70)
    print("STEP 1: Full-shard inventory (real data, actually decoded)")
    print("=" * 70)
    records, corrupted, pf = decode_full_shard_metadata(SHARD_PATH)
    summary = summarize(records)
    print(json.dumps(summary, indent=2))

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    with open(OUT_DIR / "full_shard_inventory_summary.json", "w") as f:
        json.dump(summary, f, indent=2)
    with open(OUT_DIR / "corrupted_rows.json", "w") as f:
        json.dump(corrupted, f, indent=2)

    print()
    print("=" * 70)
    print("STEP 2: Stratified development sample (target 300, seed=42)")
    print("=" * 70)
    for r in records:
        r["duration_sec_for_strat"] = r.get("duration_sec")
    strat_input = [
        {**r, "duration_sec": r.get("duration_sec") or 0.0}
        for r in records
    ]
    sample = stratified_reservoir_sample(strat_input, target_n=300, seed=42)
    print(f"Sampled {len(sample)} rows (target 300) from {len(records)} available.")

    sample_summary = summarize(sample)
    print(json.dumps(sample_summary, indent=2))
    with open(OUT_DIR / "dev_sample_summary.json", "w") as f:
        json.dump(sample_summary, f, indent=2)

    print()
    print("=" * 70)
    print("STEP 3: Materializing sampled audio as real WAV files")
    print("=" * 70)

    pf2 = read_footer_metadata(SHARD_PATH)  # re-read for column offsets
    audio_dir = OUT_DIR / "audio"
    audio_dir.mkdir(parents=True, exist_ok=True)

    # Group sample by row_group to minimize re-reading column chunks
    by_rg: dict[int, list[dict]] = {}
    for r in sample:
        by_rg.setdefault(r["_row_group"], []).append(r)

    metadata_rows = []
    n_written = 0
    n_ffmpeg_failed = 0
    for rg_idx, recs in sorted(by_rg.items()):
        rg = pf2.row_groups[rg_idx]
        audio_col = rg.columns["audio.bytes"]
        audio_mdl = pf2.max_def_levels["audio.bytes"]
        audio_vals = read_column_values(SHARD_PATH, audio_col, max_rows=rg.num_rows, max_def_level=audio_mdl)

        for r in recs:
            idx = r["_index_in_group"]
            audio_bytes = audio_vals[idx]
            if audio_bytes is None:
                continue
            out_path = audio_dir / f"{r['id']}.wav"
            proc = subprocess.run(
                ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
                 "-i", "pipe:0", "-ar", "16000", "-ac", "1", str(out_path)],
                input=audio_bytes, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
            )
            if proc.returncode != 0 or not out_path.exists():
                n_ffmpeg_failed += 1
                continue
            n_written += 1
            metadata_rows.append({
                "id": r["id"],
                "audio_path": str(out_path),
                "language": r["language"],
                "endpoint_bool": r["endpoint_bool"],
                "midfiller": r["midfiller"],
                "endfiller": r["endfiller"],
                "synthetic": r["synthetic"],
                "dataset": r["dataset"],
                "duration_sec": r.get("duration_sec"),
            })

    print(f"Materialized {n_written} real WAV files ({n_ffmpeg_failed} ffmpeg failures).")

    import csv
    with open(OUT_DIR / "metadata.csv", "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=["id", "audio_path", "language", "endpoint_bool",
                                                  "midfiller", "endfiller", "synthetic", "dataset", "duration_sec"])
        writer.writeheader()
        for row in metadata_rows:
            writer.writerow(row)

    print(f"Wrote metadata for {len(metadata_rows)} rows to {OUT_DIR / 'metadata.csv'}")


if __name__ == "__main__":
    main()