File size: 5,028 Bytes
f1ef7e2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""
Build a North-American-accent test set from apptek_callcenter_dialogues.

Downloads the parquet index, filters to the requested accents (all domains),
pairs channel1 (agent) + channel2 (customer) into calls, downloads the per-channel
WAVs, and writes a manifest.json the batch runner / evaluator consume.

Resumable: skips WAVs already on disk.

Usage:
  python build_na_testset.py --accents en-CA en-US_General
  python build_na_testset.py --accents en-CA en-US_General --limit-per-accent 3   # pilot
"""

import os
import re
import json
import time
import argparse
import urllib.request
from collections import defaultdict

import pyarrow.parquet as pq

REPO       = "apptek-com/apptek_callcenter_dialogues"
PARQUET_URL = (
    "https://huggingface.co/datasets/apptek-com/apptek_callcenter_dialogues"
    "/resolve/refs%2Fconvert%2Fparquet/default/test/0000.parquet"
)
OUT_DIR    = r"d:\Desktop\ai-ml-capstone\data\na_testset"
INDEX_PARQUET = os.path.join(OUT_DIR, "_index.parquet")
MANIFEST   = os.path.join(OUT_DIR, "manifest.json")
UA         = {"User-Agent": "Mozilla/5.0"}


def fetch(url, dest, retries=3):
    for attempt in range(1, retries + 1):
        try:
            req = urllib.request.Request(url, headers=UA)
            with urllib.request.urlopen(req, timeout=120) as r, open(dest, "wb") as f:
                f.write(r.read())
            return True
        except Exception as e:
            print(f"      attempt {attempt}/{retries} failed: {e}")
            time.sleep(2 * attempt)
    return False


def path_to_url(hf_path):
    # hf://datasets/<repo>@<rev>/test/<accent>/audio/<file>.wav
    m = re.search(r"@([0-9a-f]+)/(.*)$", hf_path)
    rev, path_in_repo = m.group(1), m.group(2)
    return f"https://huggingface.co/datasets/{REPO}/resolve/{rev}/{path_in_repo}"


def call_id_of(hf_path):
    fname = hf_path.split("/")[-1]                 # en_CA_Agriculture_1586885_channel1.wav
    return re.sub(r"_channel[12]\.wav$", "", fname)  # en_CA_Agriculture_1586885


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--accents", nargs="+", default=["en-CA", "en-US_General"])
    ap.add_argument("--limit-per-accent", type=int, default=None,
                    help="cap calls per accent (for a pilot run)")
    args = ap.parse_args()

    os.makedirs(OUT_DIR, exist_ok=True)

    # 1. Index
    if not os.path.exists(INDEX_PARQUET):
        print("Downloading parquet index...")
        if not fetch(PARQUET_URL, INDEX_PARQUET):
            raise SystemExit("Failed to download parquet index")
    df = pq.read_table(INDEX_PARQUET).to_pandas()
    df["path"] = df["audio"].apply(lambda a: a["path"])

    # 2. Filter + pair channels into calls
    manifest = []
    for accent in args.accents:
        sub = df[df["accent"] == accent]
        calls = defaultdict(dict)
        for _, row in sub.iterrows():
            cid = call_id_of(row["path"])
            ch  = "agent" if row["path"].endswith("channel1.wav") else "customer"
            calls[cid][ch] = row

        complete = {cid: v for cid, v in calls.items() if "agent" in v and "customer" in v}
        cids = sorted(complete)
        if args.limit_per_accent:
            cids = cids[: args.limit_per_accent]
        print(f"\n{accent}: {len(complete)} complete calls"
              f"{f' (using first {len(cids)})' if args.limit_per_accent else ''}")

        adir = os.path.join(OUT_DIR, accent)
        os.makedirs(adir, exist_ok=True)

        for i, cid in enumerate(cids, 1):
            a_row, c_row = complete[cid]["agent"], complete[cid]["customer"]
            a_wav = os.path.join(adir, f"{cid}_agent.wav")
            c_wav = os.path.join(adir, f"{cid}_customer.wav")

            for wav, row in [(a_wav, a_row), (c_wav, c_row)]:
                if os.path.exists(wav) and os.path.getsize(wav) > 0:
                    continue
                print(f"  [{i}/{len(cids)}] {os.path.basename(wav)}")
                if not fetch(path_to_url(row["path"]), wav):
                    print(f"      SKIP (download failed)")

            if os.path.exists(a_wav) and os.path.exists(c_wav):
                manifest.append({
                    "call_id": cid,
                    "accent":  accent,
                    "domain":  a_row["domain"],
                    "agent_wav":    os.path.relpath(a_wav, OUT_DIR),
                    "customer_wav": os.path.relpath(c_wav, OUT_DIR),
                    "agent_gender":    a_row["gender"],
                    "customer_gender": c_row["gender"],
                    "agent_transcript":    a_row["text"],
                    "customer_transcript": c_row["text"],
                })

    with open(MANIFEST, "w", encoding="utf-8") as f:
        json.dump(manifest, f, indent=2)

    by_accent = defaultdict(int)
    for m in manifest:
        by_accent[m["accent"]] += 1
    print(f"\nManifest written: {MANIFEST}")
    print(f"Total calls: {len(manifest)}  " + "  ".join(f"{k}={v}" for k, v in by_accent.items()))


if __name__ == "__main__":
    main()