Spaces:
Running
Running
| import os | |
| import json | |
| import time | |
| import socket | |
| import threading | |
| import requests | |
| import pyarrow.parquet as pq | |
| import gc | |
| from pathlib import Path | |
| from huggingface_hub import HfApi | |
| # ── Config ─────────────────────────────────────────────────────────────────── | |
| HF_TOKEN = os.environ.get("HF_TOKEN") | |
| RAW_DIR = "/data/raw" | |
| STATE_FILE = "/data/state.json" | |
| WORKER_TIMEOUT = 700 | |
| MAX_BUFFERED = 999999 | |
| os.makedirs(RAW_DIR, exist_ok=True) | |
| api = HfApi(token=HF_TOKEN) | |
| AUTH_HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"} | |
| # ── Sources ─────────────────────────────────────────────────────────────────── | |
| SOURCES = [ | |
| { | |
| "name" : "fineweb", | |
| "type" : "hf_list", | |
| "repo" : "HuggingFaceFW/fineweb-edu", | |
| "prefix" : "data/CC-MAIN-2025-26", | |
| "skip" : 5, | |
| "take" : 10, | |
| "text_col": "text", | |
| }, | |
| { | |
| "name" : "wikipedia", | |
| "type" : "hf_list", | |
| "repo" : "wikimedia/wikipedia", | |
| "prefix" : "20231101.en/train-", | |
| "skip" : 2, | |
| "take" : 18, | |
| "text_col": "text", | |
| }, | |
| { | |
| "name" : "openwebmath", | |
| "type" : "hf_list", | |
| "repo" : "open-web-math/open-web-math", | |
| "prefix" : "data/train-", | |
| "skip" : 0, | |
| "take" : 6, | |
| "text_col": "text", | |
| }, | |
| { | |
| "name" : "code", | |
| "type" : "url_list", | |
| "text_col": "text", | |
| "fmt" : "jsonl", | |
| "urls" : [ | |
| f"https://huggingface.co/buckets/Neon-tech/Dataset-arranger/resolve/by-language/{lang}/shard_{str(i).zfill(6)}.jsonl?download=true" | |
| for lang in ["C", "C++", "Java", "Go", "Rust", "Ruby", "PHP", "SQL", "C#", "Scala", "Lua", "Perl", "CSS"] | |
| for i in range(2) | |
| ], | |
| }, | |
| ] | |
| # ── Keep-alive ──────────────────────────────────────────────────────────────── | |
| def serve(): | |
| s = socket.socket(socket.AF_INET, socket.SOCK_STREAM) | |
| s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) | |
| s.bind(("0.0.0.0", 7860)) | |
| s.listen(5) | |
| print("✓ Listening on port 7860") | |
| while True: | |
| conn, _ = s.accept() | |
| conn.send(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nOK") | |
| conn.close() | |
| # ── State ───────────────────────────────────────────────────────────────────── | |
| def load_state(): | |
| if os.path.exists(STATE_FILE): | |
| with open(STATE_FILE) as f: | |
| state = json.load(f) | |
| shards = state["shards"] | |
| queue = state.get("queue", []) | |
| done = sum(1 for v in shards.values() if v["status"] == "done") | |
| claimed = sum(1 for v in shards.values() if v["status"] == "claimed") | |
| pending = sum(1 for v in shards.values() if v["status"] == "pending") | |
| print(f"Resuming — {done} done / {claimed} claimed / {pending} buffered / {len(queue)} queued") | |
| else: | |
| state = {"shards": {}, "queue": []} | |
| print("Starting fresh") | |
| return state | |
| def save_state(state): | |
| tmp = STATE_FILE + ".tmp" | |
| with open(tmp, "w") as f: | |
| json.dump(state, f, indent=2) | |
| os.replace(tmp, STATE_FILE) | |
| # ── Discover ────────────────────────────────────────────────────────────────── | |
| def discover_all(state): | |
| known_urls = {v["url"] for v in state["shards"].values()} | {e["url"] for e in state.get("queue", [])} | |
| new_count = 0 | |
| for src in SOURCES: | |
| name = src["name"] | |
| print(f"\nDiscovering: {name}") | |
| if src["type"] == "hf_list": | |
| all_files = sorted([ | |
| f for f in api.list_repo_files(src["repo"], repo_type="dataset") | |
| if f.startswith(src["prefix"]) and f.endswith(".parquet") | |
| ]) | |
| selected = all_files[src["skip"]: src["skip"] + src["take"]] | |
| base_url = f"https://huggingface.co/datasets/{src['repo']}/resolve/main/" | |
| urls = [base_url + f for f in selected] | |
| fmt = "parquet" | |
| else: | |
| urls = src["urls"] | |
| fmt = src.get("fmt", "parquet") | |
| added = 0 | |
| for url in urls: | |
| if url not in known_urls: | |
| state["queue"].append({ | |
| "url" : url, | |
| "source" : name, | |
| "text_col" : src["text_col"], | |
| "fmt" : fmt, | |
| }) | |
| known_urls.add(url) | |
| new_count += 1 | |
| added += 1 | |
| print(f" {name}: {len(urls)} files | {added} new added to queue") | |
| save_state(state) | |
| print(f"\nTotal queued: {len(state['queue'])} | In state: {len(state['shards'])}") | |
| # ── Reclaim stale ───────────────────────────────────────────────────────────── | |
| def reclaim_stale(state): | |
| now = time.time() | |
| reclaimed = 0 | |
| for name, info in state["shards"].items(): | |
| if info["status"] == "claimed" and info.get("claimed_at"): | |
| if now - info["claimed_at"] > WORKER_TIMEOUT: | |
| print(f" ⚠ Reclaiming: {name}") | |
| info["status"] = "pending" | |
| info["worker"] = None | |
| info["claimed_at"] = None | |
| reclaimed += 1 | |
| if reclaimed: | |
| save_state(state) | |
| # ── Parquet → JSONL ─────────────────────────────────────────────────────────── | |
| def parquet_to_jsonl(parquet_path, jsonl_path, text_col): | |
| """Stream parquet batch by batch → write one JSON line per doc. No full load.""" | |
| pf = pq.ParquetFile(parquet_path) | |
| n_written = 0 | |
| with open(jsonl_path, "w", encoding="utf-8") as out: | |
| for batch in pf.iter_batches(batch_size=1_000, columns=[text_col]): | |
| texts = batch.column(text_col).to_pylist() | |
| for text in texts: | |
| if text and isinstance(text, str) and text.strip(): | |
| out.write(json.dumps({"text": text.strip()}, ensure_ascii=False) + "\n") | |
| n_written += 1 | |
| del texts | |
| gc.collect() | |
| return n_written | |
| # ── Download loop ───────────────────────────────────────────────────────────── | |
| def download_loop(state): | |
| while True: | |
| try: | |
| with open(STATE_FILE) as f: | |
| fresh = json.load(f) | |
| state["shards"] = fresh["shards"] | |
| state["queue"] = fresh.get("queue", []) | |
| except Exception: | |
| pass | |
| reclaim_stale(state) | |
| buffered = sum(1 for v in state["shards"].values() if v["status"] == "pending") | |
| if buffered >= MAX_BUFFERED: | |
| time.sleep(30) | |
| continue | |
| if not state["queue"]: | |
| done = sum(1 for v in state["shards"].values() if v["status"] == "done") | |
| total = len(state["shards"]) | |
| if done == total and total > 0: | |
| print("✓ All shards complete!") | |
| break | |
| print(" Queue empty — sleeping...") | |
| time.sleep(60) | |
| continue | |
| entry = state["queue"][0] | |
| url = entry["url"] | |
| source = entry["source"] | |
| text_col = entry["text_col"] | |
| fmt = entry.get("fmt", "parquet") | |
| lang = url.split("?")[0].split("/")[-2] | |
| base_name = url.split("?")[0].split("/")[-1].replace(".parquet", "").replace(".jsonl", "") | |
| shard_name = f"{source}__{base_name}_{lang}.jsonl" | |
| jsonl_path = Path(RAW_DIR) / shard_name | |
| tmp_path = Path(RAW_DIR) / f"{shard_name}.tmp" | |
| print(f" Downloading: {source} | {base_name}") | |
| try: | |
| resp = requests.get(url, headers=AUTH_HEADERS, timeout=300, stream=True) | |
| resp.raise_for_status() | |
| with open(tmp_path, "wb") as f: | |
| for chunk in resp.iter_content(chunk_size=8 * 1024 * 1024): | |
| f.write(chunk) | |
| except Exception as e: | |
| print(f" ✗ Download failed: {e} — retrying in 30s") | |
| tmp_path.unlink(missing_ok=True) | |
| time.sleep(30) | |
| continue | |
| if fmt == "parquet": | |
| print(f" Converting → jsonl: {shard_name}") | |
| try: | |
| n = parquet_to_jsonl(tmp_path, jsonl_path, text_col) | |
| tmp_path.unlink(missing_ok=True) | |
| print(f" ✓ {n:,} docs") | |
| except Exception as e: | |
| print(f" ✗ Convert failed: {e}") | |
| tmp_path.unlink(missing_ok=True) | |
| jsonl_path.unlink(missing_ok=True) | |
| time.sleep(30) | |
| continue | |
| else: | |
| tmp_path.rename(jsonl_path) | |
| state["queue"].pop(0) | |
| state["shards"][shard_name] = { | |
| "status" : "pending", | |
| "url" : url, | |
| "source" : source, | |
| "worker" : None, | |
| "claimed_at": None, | |
| "error" : None, | |
| } | |
| save_state(state) | |
| print(f" ✓ Ready: {shard_name}") | |
| time.sleep(3) | |
| # ── Monitor ─────────────────────────────────────────────────────────────────── | |
| def monitor_loop(): | |
| while True: | |
| time.sleep(120) | |
| try: | |
| with open(STATE_FILE) as f: | |
| s = json.load(f) | |
| shards = s["shards"] | |
| queue = s.get("queue", []) | |
| done = sum(1 for v in shards.values() if v["status"] == "done") | |
| claimed = sum(1 for v in shards.values() if v["status"] == "claimed") | |
| pending = sum(1 for v in shards.values() if v["status"] == "pending") | |
| total = len(shards) + len(queue) | |
| pct = (done / total * 100) if total else 0 | |
| src_done = {} | |
| for v in shards.values(): | |
| src = v.get("source", "?") | |
| if v["status"] == "done": | |
| src_done[src] = src_done.get(src, 0) + 1 | |
| print(f"[MONITOR] {done}/{total} ({pct:.1f}%) | {claimed} active | {pending} buffered | {len(queue)} queued") | |
| for src, cnt in sorted(src_done.items()): | |
| print(f" {src}: {cnt} done") | |
| except Exception: | |
| pass | |
| # ── Entry point ─────────────────────────────────────────────────────────────── | |
| if __name__ == "__main__": | |
| threading.Thread(target=serve, daemon=True).start() | |
| state = load_state() | |
| discover_all(state) | |
| threading.Thread(target=monitor_loop, daemon=True).start() | |
| threading.Thread(target=download_loop, args=(state,), daemon=True).start() | |
| while True: | |
| time.sleep(60) |