Update app.py

app.py

@@ -1,7 +1,202 @@
-import shutil
+import os
+import json
+import time
+import socket
+import threading
+import io
+import requests
+import pandas as pd
 from pathlib import Path
+from tokenizers import Tokenizer
+from huggingface_hub import HfApi
 
-for folder in ["/data/by-language", "/data/progress_state.json"]:
-    p = Path(folder)
-    if p.exists():
-        shutil.rmtree(p)
+# ── Config ───────────────────────────────────────────────────────────────────
+HF_TOKEN       = os.environ.get("HF_TOKEN")
+DATASET_REPO   = "Neon-coding/github-code-raw"
+TOK_PATH       = "/data/tokenizer.json"
+OUT_DIR        = "/data/by-language"
+STATE_FILE     = "/data/progress_state.json"
+TOTAL_PARQUETS = 880
+SHARD_TOKENS   = 50_000_000  # 50M tokens per shard
+
+PARQUET_URL = (
+    "https://huggingface.co/datasets/codeparrot/github-code-clean"
+    "/resolve/main/data/train-{i:05d}-of-00880.parquet"
+)
+
+os.makedirs(OUT_DIR, exist_ok=True)
+
+api = HfApi(token=HF_TOKEN)
+
+# ── Port 7860 — keeps Space green ────────────────────────────────────────────
+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)
+        print(f"Resuming — {len(state['done'])} / {TOTAL_PARQUETS} parquets done")
+    else:
+        state = {
+            "done":        [],
+            "lang_shards": {},
+            "lang_tokens": {},
+        }
+        print("Starting fresh")
+    return state
+
+def save_state(state, retries=3, delay=5):
+    for attempt in range(retries):
+        try:
+            with open(STATE_FILE, "w") as f:
+                json.dump(state, f, indent=2)
+            return
+        except OSError as e:
+            print(f"  ⚠ State save attempt {attempt + 1} failed: {e}")
+            if attempt < retries - 1:
+                time.sleep(delay)
+    print("  ✗ State save failed after all retries — continuing")
+
+# ── Shard buffers — global per language, persist across parquets ─────────────
+buffers = {}
+
+def get_buffer(lang):
+    if lang not in buffers:
+        buffers[lang] = {"rows": [], "token_count": 0}
+    return buffers[lang]
+
+def flush_shard(lang, rows, state):
+    shard_idx  = state["lang_shards"].get(lang, 0)
+    lang_dir   = Path(OUT_DIR) / lang
+    lang_dir.mkdir(parents=True, exist_ok=True)
+    shard_name = f"shard_{shard_idx:06d}.jsonl"
+    shard_path = lang_dir / shard_name
+
+    with open(shard_path, "w", encoding="utf-8") as f:
+        for row in rows:
+            f.write(json.dumps(row, ensure_ascii=False) + "\n")
+
+    tok_in_shard = sum(r["token_count"] for r in rows)
+    state["lang_shards"][lang] = shard_idx + 1
+    state["lang_tokens"][lang] = state["lang_tokens"].get(lang, 0) + tok_in_shard
+    print(f"  ✓ {lang}/{shard_name} | {len(rows)} samples | {tok_in_shard:,} tokens")
+
+# ── Main processing loop ─────────────────────────────────────────────────────
+def process(tokenizer, state):
+    for i in range(TOTAL_PARQUETS):
+        if i in state["done"]:
+            print(f"[{i:06d}/{TOTAL_PARQUETS}] SKIP")
+            continue
+
+        url = PARQUET_URL.format(i=i)
+        print(f"[{i:06d}/{TOTAL_PARQUETS}] Downloading...")
+
+        try:
+            resp = requests.get(
+                url,
+                headers={"Authorization": f"Bearer {HF_TOKEN}"},
+                timeout=180,
+            )
+            resp.raise_for_status()
+            df = pd.read_parquet(io.BytesIO(resp.content))
+        except Exception as e:
+            print(f"[{i:06d}] Download error: {e} — skipping")
+            continue
+
+        print(f"[{i:06d}] {len(df):,} rows | {df['language'].nunique()} languages")
+
+        # row by row — constant memory
+        for row_tuple in df.itertuples(index=False):
+            lang        = row_tuple.language
+            text        = row_tuple.code if row_tuple.code else ""
+            repo        = row_tuple.repo_name
+            fpath       = row_tuple.path
+            lic         = row_tuple.license
+
+            if not text.strip():
+                continue
+
+            enc         = tokenizer.encode(text)
+            token_count = len(enc.ids)
+
+            if token_count < 2:
+                continue
+
+            buf = get_buffer(lang)
+            row = {
+                "text":        text,
+                "token_count": token_count,
+                "repo":        repo,
+                "path":        fpath,
+                "license":     lic,
+            }
+
+            if buf["token_count"] + token_count > SHARD_TOKENS and buf["rows"]:
+                flush_shard(lang, buf["rows"], state)
+                save_state(state)
+                buf["rows"]        = []
+                buf["token_count"] = 0
+
+            buf["rows"].append(row)
+            buf["token_count"] += token_count
+
+        del df
+
+        state["done"].append(i)
+        save_state(state)
+        print(f"[{i:06d}] ✓ Complete")
+
+    # ── Flush remaining partial shards ────────────────────────────────────────
+    print("\nFlushing remaining buffers...")
+    for lang, buf in buffers.items():
+        if buf["rows"]:
+            flush_shard(lang, buf["rows"], state)
+    save_state(state)
+
+    # ── Write meta.json per language ──────────────────────────────────────────
+    print("\nWriting meta.json per language...")
+    for lang in state["lang_tokens"]:
+        meta = {
+            "language":     lang,
+            "total_tokens": state["lang_tokens"][lang],
+            "total_shards": state["lang_shards"].get(lang, 0),
+        }
+        meta_path = Path(OUT_DIR) / lang / "meta.json"
+        with open(meta_path, "w") as f:
+            json.dump(meta, f, indent=2)
+        print(f"  {lang}: {meta['total_tokens']:,} tokens | {meta['total_shards']} shards")
+
+    # ── Push everything to HF dataset repo ───────────────────────────────────
+    print(f"\nPushing to {DATASET_REPO}...")
+    api.upload_folder(
+        folder_path=OUT_DIR,
+        repo_id=DATASET_REPO,
+        repo_type="dataset",
+        token=HF_TOKEN,
+    )
+    print("\n✓ All done!")
+
+# ── Entry point ──────────────────────────────────────────────────────────────
+if __name__ == "__main__":
+    threading.Thread(target=serve, daemon=True).start()
+
+    print("✓ Loading tokenizer from /data/tokenizer.json...")
+    tokenizer = Tokenizer.from_file(TOK_PATH)
+    print(f"✓ Tokenizer loaded | vocab: {tokenizer.get_vocab_size():,}")
+
+    state = load_state()
+
+    threading.Thread(target=process, args=(tokenizer, state), daemon=True).start()
+
+    while True:
+        time.sleep(60)