| |
| import io |
| import re |
| from pathlib import Path |
|
|
| import requests |
| import pyarrow.parquet as pq |
|
|
| BASE = "https://huggingface.co/datasets/Estwld/empathetic_dialogues_llm/resolve/refs%2Fconvert%2Fparquet/default" |
| SPLITS = {"train": "train", "val": "valid"} |
| OUT_DIR = Path("data/sft") |
|
|
| BOS = "<bos>" |
| EOS = "<eos>" |
| ROLES = ("<|user|>", "<|assistant|>") |
|
|
| _MULTI = re.compile(r"\s+") |
| _APOS = re.compile(r"\s+'\s*") |
| _PUNCT = re.compile(r"\s+([.,!?;:])") |
| _CODE = re.compile(r"`|</|/>|def |import |class |printf|console\.log|#include|<tool") |
|
|
|
|
| def clean_utt(u): |
| u = u.replace("’", "'").replace("‘", "'").replace("“", '"').replace("”", '"') |
| u = _MULTI.sub(" ", u.strip()) |
| u = _APOS.sub("'", u) |
| u = _PUNCT.sub(r"\1", u) |
| return _MULTI.sub(" ", u).strip() |
|
|
|
|
| def has_code(turns): |
| return bool(_CODE.search(" ".join(turns))) |
|
|
|
|
| def fetch_conversations(split): |
| url = f"{BASE}/{split}/0000.parquet" |
| raw = requests.get(url, timeout=120).content |
| table = pq.read_table(io.BytesIO(raw)) |
| return table.column("conversations").to_pylist() |
|
|
|
|
| def format_dialogue(turns): |
| parts = [BOS] |
| for i, text in enumerate(turns): |
| parts.append(ROLES[i % 2]) |
| parts.append(text) |
| parts.append(EOS) |
| return "".join(parts) |
|
|
|
|
| def main(): |
| OUT_DIR.mkdir(parents=True, exist_ok=True) |
| for name, split in SPLITS.items(): |
| seen = kept = dropped_short = dropped_code = 0 |
| out = OUT_DIR / f"{name}.txt" |
| with out.open("w", encoding="utf-8") as f: |
| for convo in fetch_conversations(split): |
| seen += 1 |
| turns = [t for t in (clean_utt(c["content"]) for c in convo) if t] |
| if len(turns) < 2: |
| dropped_short += 1 |
| continue |
| if has_code(turns): |
| dropped_code += 1 |
| continue |
| f.write(format_dialogue(turns) + "\n") |
| kept += 1 |
| print(f"{name}: seen={seen} kept={kept} dropped_short={dropped_short} dropped_code={dropped_code} -> {out}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
