""" Does your test set look like your users? Steps 3 to 6. ======================================================== Data: WildChat-1M (allenai/WildChat-1M, ODC-BY), real conversations people had with a free ChatGPT-based chatbot in 2023, published by the Allen Institute for AI. This script uses the first of its 14 files: 59,857 conversations from 9 April to 4 May 2023. Zhao et al., 2024, "WildChat: 1M ChatGPT Interaction Logs in the Wild." Compared with MT-Bench's 80 questions (LMSYS, FastChat). Setup (once): pip install pandas pyarrow Run: python traffic_lesson.py (profile of all traffic vs MT-Bench) python traffic_lesson.py profile Russian (profile of one language) python traffic_lesson.py sample 300 (build a 300-item test set) The first run downloads about 230 MB into ./data. """ import os import sys import urllib.request import pandas as pd FILES = { 'data/wildchat_00000.parquet': 'https://huggingface.co/datasets/allenai/WildChat-1M/resolve/main/data/train-00000-of-00014.parquet', 'data/mt_bench_questions.jsonl': 'https://huggingface.co/spaces/lmsys/mt-bench/resolve/main/data/mt_bench/question.jsonl', } CODE = r'```|\bdef |\bfunction\b|\bimport |#include|public class|console\.log|SELECT .* FROM|Traceback' def load(): os.makedirs('data', exist_ok=True) for local, url in FILES.items(): if not os.path.exists(local): print('downloading', url.split('/')[-1], '...') urllib.request.urlretrieve(url, local) d = pd.read_parquet('data/wildchat_00000.parquet', columns=['conversation', 'turn', 'language', 'toxic']) d = d[~d.toxic].copy() d['first'] = d.conversation.map(lambda c: c[0]['content']) d['words'] = d['first'].str.split().str.len() d['chars'] = d['first'].str.len() # Chinese and Japanese don't put spaces between words, so measure them in characters. nospace = d.language.isin(['Chinese', 'Japanese']) d['size'] = d['words'].where(~nospace, d['chars']) d['opening'] = d['first'].str.strip().str[:60] d['template'] = d['opening'].map(d['opening'].value_counts()) >= 20 # same opening as 20+ others d['code'] = d['first'].str.contains(CODE, regex=True) d['length'] = pd.cut(d['size'], [-1, 5, 300, 10**9], labels=['short', 'normal', 'long']) mt = pd.read_json('data/mt_bench_questions.jsonl', lines=True) mt['words'] = mt.turns.map(lambda t: len(t[0].split())) return d.drop(columns=['conversation']), mt def pct(x): return f"{x:.0%}" if x >= 0.01 or x == 0 else "<1%" def profile(d, mt, language=None): if language: d = d[d.language == language] print(f"\n{language} conversations only ({len(d):,})") else: print(f"\nAll conversations ({len(d):,}) compared with MT-Bench's {len(mt)} questions") print("\nLanguage of the conversation") top = d.language.value_counts(normalize=True).head(6) for lang, s in top.items(): print(f" {lang:<12} {pct(s):>5}") print(f" {'all others':<12} {pct(1 - top.sum()):>5} MT-Bench: 100% English") print(f"\n {'One-message conversations':<36} {pct((d.turn == 1).mean()):>10} MT-Bench: always 2 turns") print(f" {'Empty first messages':<36} {pct((d.chars == 0).mean()):>10}") d = d[d.language == 'English'] print(f"\nThe rest compares English conversations ({len(d):,}) with MT-Bench, which is all English.") unit = 'characters' if language in ('Chinese', 'Japanese') else 'words' size = d['size'] print(f"\n {'First message length, in ' + unit:<36} {'real users':>10} {'MT-Bench (words)':>16}") for q, label in [(.1, '10% are shorter than'), (.5, 'half are shorter than'), (.9, '90% are shorter than')]: print(f" {label:<36} {size.quantile(q):>10.0f} {mt.words.quantile(q):>16.0f}") print(f" {'5 ' + unit + ' or fewer':<36} {pct((size <= 5).mean()):>10} {pct((mt.words <= 5).mean()):>16}") print(f" {'over 300 ' + unit:<36} {pct((size > 300).mean()):>10} {pct((mt.words > 300).mean()):>16}") if language: print(f"\n {'One-message conversations':<36} {pct((d.turn == 1).mean()):>10} MT-Bench: always 2 turns") print(f" {'Messages containing code or errors':<36} {pct(d.code.mean()):>10}") print(f" {'Messages sharing an opening with 20+':<36} {pct(d.template.mean()):>10}") print("\nMost repeated openings") for opening, n in d.opening.value_counts().head(8).items(): shown = opening.replace(chr(10), ' ')[:60] if opening else '(empty message)' print(f" {n:>5} {shown}") def sample(d, n, floor=15, seed=0): """Give every slice at least `floor` items, then share the rest of the budget in proportion to traffic.""" d = d.copy() lang = d.language.where(d.language.isin(['English', 'Chinese', 'Russian']), 'other') d['slice'] = lang + ' / ' + d.length.astype(str) share = d.slice.value_counts(normalize=True) available = d.slice.value_counts() random_plan = (share * n).round().astype(int) base = available.clip(upper=floor) # a slice can't get more items than it has rest = n - base.sum() extra = (share * rest).round().astype(int) plan = (base + extra).clip(upper=available) print(f"\nA {n}-item test set drawn from real traffic (every slice gets at least {floor} items)") print(f" {'slice':<22} {'share of traffic':>17} {'simple random':>14} {'with a floor':>13}") for s in share.index: print(f" {s:<22} {pct(share[s]):>17} {random_plan[s]:>14} {plan[s]:>13}") print(f" {'total':<22} {'':>17} {random_plan.sum():>14} {plan.sum():>13}") rows = [d[d.slice == s].sample(plan[s], random_state=seed) for s in share.index] out = pd.concat(rows) out['weight'] = out.slice.map(share / out.slice.value_counts(normalize=True)).round(3) out[['slice', 'weight', 'first']].to_csv('test_set.csv', index=False) print(f"\nWrote test_set.csv ({len(out)} items). When you report an overall score, weight each item by its") print("'weight' column, so slices you over-sampled don't count more than they do in real traffic.") if __name__ == '__main__': d, mt = load() if len(sys.argv) >= 3 and sys.argv[1] == 'profile': profile(d, mt, sys.argv[2]) elif len(sys.argv) >= 3 and sys.argv[1] == 'sample': sample(d, int(sys.argv[2])) else: profile(d, mt)