| """Download human-written coding text into data/prompts.txt and data/replies.txt. |
| |
| prompts.txt: what people type to an assistant (OASST user turns, MBPP, CoNaLa). |
| replies.txt: OASST assistant turns (human-written replies). |
| wildchat.txt: English user turns from WildChat-1M coding conversations (never the |
| assistant turns, which are model-generated). |
| |
| Sources (all human-written, permissively licensed; see data/SOURCES.md): |
| - OpenAssistant/oasst2 (Apache-2.0) English prompter and assistant messages from |
| conversations whose opening prompt is about code |
| - google-research-datasets/mbpp (CC-BY-4.0) crowd-written task descriptions |
| - neulab/conala (MIT) curated programming intents (paired train + test) |
| |
| Only prose is kept: code blocks, code-looking lines and markdown markup are |
| dropped, and each message is cut to its first sentences (prompts: 2, replies: 5). |
| """ |
|
|
| import gzip |
| import json |
| import re |
| from pathlib import Path |
|
|
| from huggingface_hub import hf_hub_download |
|
|
| DATA = Path(__file__).resolve().parent.parent / "data" |
| MAX_WORDS = 24 |
| MAX_PROMPT_SENTENCES = 2 |
| MAX_REPLY_SENTENCES = 5 |
|
|
| |
| CODE_RE = re.compile( |
| r"\b(python|javascript|typescript|java|c\+\+|c#|rust|golang|kotlin|swift|php|ruby|perl|haskell|scala|" |
| r"sql|mysql|postgres(ql)?|sqlite|mongodb|bash|powershell|shell script|html|css|react|vue|angular|node\.?js|" |
| r"django|flask|numpy|pandas|pytorch|tensorflow|regex|regular expression|json|yaml|xml|csv|api|rest api|" |
| r"git|github|docker|kubernetes|linux|terminal|command line|compiler|compile|debug|debugging|stack trace|" |
| r"exception|segfault|code|coding|program|programming|programmer|script|function|algorithm|recursion|" |
| r"recursive|data structure|linked list|binary tree|hash ?map|array|dictionary|variable|loop|class|" |
| r"object[- ]oriented|unit test|runtime|syntax|library|framework|database|backend|frontend|webpage|website)\b", |
| re.I, |
| ) |
| |
| |
| SECRET_RE = re.compile( |
| r"sk-(?:proj-|svcacct-|admin-|ant-)?[A-Za-z0-9_-]{20,}|(?:ghp|gho|ghu|ghs|ghr)_[A-Za-z0-9]{30,}|github_pat_\w{30,}" |
| r"|(?:AKIA|ASIA)[0-9A-Z]{16}|AIza[0-9A-Za-z_-]{35}|xox[baprs]-[A-Za-z0-9-]{10,}|(?:sk|rk)_(?:live|test)_\w{20,}" |
| r"|hf_[A-Za-z0-9]{30,}|eyJ[\w-]{10,}\.[\w-]{10,}\.[\w-]{10,}|-----BEGIN [A-Z ]*PRIVATE KEY-----|\b\d{8,10}:[\w-]{35}\b" |
| ) |
| RANDOM_TOKEN_RE = re.compile(r"[A-Za-z0-9_\-+/=]{32,}") |
|
|
|
|
| def looks_secret(text: str) -> bool: |
| if SECRET_RE.search(text): |
| return True |
| return any(re.search(r"[A-Za-z]", t) and re.search(r"\d", t) for t in RANDOM_TOKEN_RE.findall(text)) |
|
|
|
|
| CODE_LINE_RE = re.compile(r"[{};=<>]|^\s*(def|class|import|from|for|if|return|#include|const|let|var)\b|\w\(.*\)") |
| SENT_SPLIT_RE = re.compile(r"(?<=[.?!])\s+") |
|
|
|
|
| def clean(text: str, max_sentences: int = MAX_PROMPT_SENTENCES) -> list[str]: |
| """Prose sentences from a message, without code or markdown.""" |
| text = re.sub(r"```.*?(```|$)", " ", text, flags=re.S) |
| lines = [] |
| for ln in text.splitlines(): |
| ln = re.sub(r"^\s*(#+|[-*+]|\d+[.)])\s+", "", ln) |
| ln = re.sub(r"\*\*|__|\[([^\]]*)\]\([^)]*\)", r"\1", ln) |
| if ln.strip() and not CODE_LINE_RE.search(ln): |
| lines.append(ln) |
| text = " ".join(" ".join(lines).split()) |
| text = re.sub(r"`([^`]{1,30})`", r"\1", text) |
| text = re.sub(r"`[^`]*`", " ", text) |
| text = re.sub(r"https?://\S+", " ", text) |
| if not text.isascii(): |
| return [] |
| out = [] |
| for sent in SENT_SPLIT_RE.split(text)[:max_sentences]: |
| words = sent.split() |
| if 3 <= len(words) <= MAX_WORDS and not looks_secret(sent): |
| out.append(" ".join(words)) |
| return out |
|
|
|
|
| def oasst() -> tuple[list[str], list[str]]: |
| path = hf_hub_download("OpenAssistant/oasst2", "2023-11-05_oasst2_ready.messages.jsonl.gz", repo_type="dataset") |
| rows = [json.loads(line) for line in gzip.open(path, "rt")] |
| rows = [ |
| r for r in rows |
| if r["lang"] == "en" and not r["deleted"] and r["review_result"] and not r["synthetic"] |
| ] |
| coding_trees = {r["message_tree_id"] for r in rows if r["parent_id"] is None and CODE_RE.search(r["text"])} |
| prompts, replies = [], [] |
| for r in rows: |
| if r["message_tree_id"] in coding_trees: |
| if r["role"] == "prompter": |
| prompts += clean(r["text"], MAX_PROMPT_SENTENCES) |
| else: |
| replies += clean(r["text"], MAX_REPLY_SENTENCES) |
| return prompts, replies |
|
|
|
|
| def mbpp() -> list[str]: |
| import pyarrow.parquet as pq |
|
|
| out = [] |
| for split in ["train", "test", "validation", "prompt"]: |
| path = hf_hub_download("google-research-datasets/mbpp", f"full/{split}-00000-of-00001.parquet", repo_type="dataset") |
| for text in pq.read_table(path, columns=["text"]).column("text").to_pylist(): |
| out += clean(text) |
| return out |
|
|
|
|
| def conala() -> list[str]: |
| out = [] |
| for split in ["train", "test"]: |
| path = hf_hub_download("neulab/conala", f"data/conala-paired-{split}.json", repo_type="dataset") |
| for line in open(path): |
| if line.strip(): |
| out += clean(json.loads(line)["intent"]) |
| return out |
|
|
|
|
| def wildchat() -> list[str]: |
| """allenai/WildChat-1M (ODC-BY): user turns of English, non-toxic coding conversations.""" |
| import pyarrow.parquet as pq |
|
|
| out = [] |
| for i in range(14): |
| path = hf_hub_download("allenai/WildChat-1M", f"data/train-{i:05d}-of-00014.parquet", repo_type="dataset") |
| table = pq.read_table(path, columns=["conversation", "language", "toxic"]) |
| convs, langs, toxic = (table.column(c).to_pylist() for c in ["conversation", "language", "toxic"]) |
| for conv, lang, tox in zip(convs, langs, toxic): |
| users = [m for m in conv if m["role"] == "user"] |
| if lang != "English" or tox or not users or not CODE_RE.search(users[0]["content"]): |
| continue |
| for m in users: |
| if not m.get("redacted") and not m.get("toxic") and m.get("language") in (None, "English"): |
| out += clean(m["content"], MAX_PROMPT_SENTENCES) |
| print(f"wildchat shard {i + 1}/14: {len(out)} sentences", flush=True) |
| return out |
|
|
|
|
| def main() -> None: |
| oasst_prompts, oasst_replies = oasst() |
| sources = { |
| "prompts": [("oasst2_user", oasst_prompts), ("mbpp", mbpp()), ("conala", conala())], |
| "replies": [("oasst2_assistant", oasst_replies)], |
| "wildchat": [("wildchat_user", wildchat())], |
| } |
| seen, report = set(), {} |
| DATA.mkdir(exist_ok=True) |
| for kind, parts in sources.items(): |
| lines = [] |
| for name, got in parts: |
| report[name] = len(got) |
| for s in got: |
| if s.lower() not in seen: |
| seen.add(s.lower()) |
| lines.append(s) |
| (DATA / f"{kind}.txt").write_text("\n".join(lines) + "\n") |
| report[f"{kind}_unique"] = len(lines) |
| print(json.dumps(report, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
