Update app.py

app.py

@@ -1,118 +1,214 @@
-# process.py
-from huggingface_hub import snapshot_download, HfApi
-from tokenizers import Tokenizer
 import os
 import json
+import time
 import threading
+import io
 import requests
 import pandas as pd
 from pathlib import Path
-from concurrent.futures import ThreadPoolExecutor
-
-# ── Config ──
-HF_USERNAME  = "Neon-coding"
-DATASET_NAME = "github-code-raw"
-RAW_DIR      = "/data/codeparrot-raw"
-OUT_DIR      = "/data/by-language"
-STATE_FILE   = "/data/progress_state.json"
-TOK_PATH     = "/data/tokenizer.json"
+from tokenizers import Tokenizer
+from huggingface_hub import HfApi
+
+# ── Config ───────────────────────────────────────────────────────────────────
+HF_TOKEN      = os.environ.get("HF_TOKEN")
+HF_USERNAME   = "Neon-coding"
+DATASET_REPO  = f"{HF_USERNAME}/github-code-raw"
+BUCKET_REPO   = f"{HF_USERNAME}/ureola-bucket"   # where tokenizer.json lives
+OUT_DIR       = "/data/by-language"
+STATE_FILE    = "/data/progress_state.json"
+TOK_FILENAME  = "tokenizer.json"
+TOTAL_PARQUETS = 880
+SHARD_TOKENS  = 100_000   # exactly 100k tokens per shard file
+
+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)
-os.makedirs(RAW_DIR, exist_ok=True)
-
-# ── Load tokenizer (already in bucket) ──
-print("✓ Loading tokenizer from bucket...")
-tokenizer = Tokenizer.from_file(TOK_PATH)
-SEP_TOKEN = tokenizer.token_to_id("<eos>")
-print(f"✓ Tokenizer loaded | vocab: {tokenizer.get_vocab_size():,}")
-
-# ── Load state ──
-if os.path.exists(STATE_FILE):
-    with open(STATE_FILE) as f:
-        state = json.load(f)
-    print(f"Resuming — {len(state['processed_files'])} files already done")
-else:
-    state = {"processed_files": [], "lang_tokens": {}}
-    print("Starting fresh")
-
-lock = threading.Lock()
-
-def save_state():
+os.makedirs("/data", exist_ok=True)
+
+api = HfApi(token=HF_TOKEN)
+
+# ── Pull tokenizer.json from bucket ─────────────────────────────────────────
+def load_tokenizer():
+    tok_path = f"/data/{TOK_FILENAME}"
+    if not os.path.exists(tok_path):
+        print("Pulling tokenizer.json from bucket...")
+        api.hf_hub_download(
+            repo_id=BUCKET_REPO,
+            repo_type="dataset",
+            filename=TOK_FILENAME,
+            local_dir="/data",
+            token=HF_TOKEN,
+        )
+    tokenizer = Tokenizer.from_file(tok_path)
+    print(f"✓ Tokenizer loaded | vocab: {tokenizer.get_vocab_size():,}")
+    return tokenizer
+
+# ── 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'])} parquets done")
+    else:
+        state = {
+            "done": [],           # list of parquet indices completed
+            "lang_shards": {},    # {lang: current shard index}
+            "lang_tokens": {},    # {lang: total tokens written so far}
+        }
+        print("Starting fresh")
+    return state
+
+def save_state(state):
     with open(STATE_FILE, "w") as f:
         json.dump(state, f, indent=2)
 
-# ── Download codeparrot ──
-print("\nDownloading codeparrot/github-code-clean...")
-local_dir = snapshot_download(
-    repo_id="codeparrot/github-code-clean",
-    repo_type="dataset",
-    local_dir=RAW_DIR,
-)
-
-parquet_files = sorted(Path(local_dir).rglob("*.parquet"))
-print(f"Found {len(parquet_files)} parquet files")
-
-# ── Process each file ──
-def process_file(path):
-    fname = str(path)
-
-    if fname in state["processed_files"]:
-        print(f"  SKIP {path.name}")
-        return
-
-    try:
-        df = pd.read_parquet(path)
-
-        for lang, group in df.groupby("language"):
-            lang_dir = os.path.join(OUT_DIR, lang)
-            os.makedirs(lang_dir, exist_ok=True)
-            out = os.path.join(lang_dir, f"{path.stem}.jsonl")
-
-            if os.path.exists(out):
-                continue
-
-            texts     = group["code"].dropna().tolist()
-            encoded   = tokenizer.encode_batch(texts)
-            tok_count = sum(len(e.ids) for e in encoded)
-
-            group[["code"]].rename(columns={"code": "text"}).to_json(
-                out, orient="records", lines=True
+# ── Shard buffer: one per language, persists across parquets ─────────────────
+# buffers[lang] = {"rows": [...], "token_count": N}
+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):
+    """Write rows to a new shard file and upload to HF dataset repo."""
+    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:05d}.jsonl"
+    shard_path = lang_dir / shard_name
+
+    with open(shard_path, "w") as f:
+        for row in rows:
+            f.write(json.dumps(row, ensure_ascii=False) + "\n")
+
+    # upload to HF
+    api.upload_file(
+        path_or_fileobj=str(shard_path),
+        path_in_repo=f"{lang}/{shard_name}",
+        repo_id=DATASET_REPO,
+        repo_type="dataset",
+        token=HF_TOKEN,
+    )
+    print(f"  ✓ Uploaded {lang}/{shard_name} | {len(rows)} samples")
+
+    # update state
+    state["lang_shards"][lang] = shard_idx + 1
+    state["lang_tokens"][lang] = state["lang_tokens"].get(lang, 0) + sum(
+        r["token_count"] for r in rows
+    )
+
+# ── Core processing loop ─────────────────────────────────────────────────────
+def process(tokenizer, state):
+    for i in range(TOTAL_PARQUETS):
+        if i in state["done"]:
+            print(f"[{i:05d}] SKIP")
+            continue
+
+        url = PARQUET_URL.format(i=i)
+        print(f"[{i:05d}] Downloading...")
+
+        try:
+            resp = requests.get(
+                url,
+                headers={"Authorization": f"Bearer {HF_TOKEN}"},
+                timeout=120,
             )
+            resp.raise_for_status()
+            df = pd.read_parquet(io.BytesIO(resp.content))
+        except Exception as e:
+            print(f"[{i:05d}] Download error: {e} — skipping")
+            continue
 
-            with lock:
-                state["lang_tokens"][lang] = state["lang_tokens"].get(lang, 0) + tok_count
-
-        with lock:
-            state["processed_files"].append(fname)
-            save_state()
-
-        print(f"  ✓ {path.name} | langs: {list(df['language'].unique())}")
-
-    except Exception as e:
-        print(f"  ✗ {path.name} ERROR: {e}")
-
+        print(f"[{i:05d}] {len(df):,} rows | processing...")
 
-with ThreadPoolExecutor(max_workers=8) as ex:
-    list(ex.map(process_file, parquet_files))
-
-# ── Save per-language meta ──
-print("\nSaving per-language meta.json...")
-for lang, total_tokens in state["lang_tokens"].items():
-    lang_dir = os.path.join(OUT_DIR, lang)
-    os.makedirs(lang_dir, exist_ok=True)
-    with open(os.path.join(lang_dir, "meta.json"), "w") as f:
-        json.dump({"language": lang, "total_tokens": total_tokens}, f, indent=2)
-    print(f"  {lang}: {total_tokens:,}")
-
-# ── Push to HF ──
-print("\nPushing to HuggingFace...")
-api = HfApi()
-api.upload_folder(
-    folder_path=OUT_DIR,
-    repo_id=f"{HF_USERNAME}/{DATASET_NAME}",
-    repo_type="dataset",
-)
-
-print("\nDone!")
-for l, t in sorted(state["lang_tokens"].items(), key=lambda x: -x[1]):
-    print(f"  {l}: {t:,}")
+        for lang, group in df.groupby("language"):
+            buf = get_buffer(lang)
+
+            texts = group["code"].fillna("").tolist()
+            repos  = group["repo_name"].tolist()
+            paths  = group["path"].tolist()
+            licenses = group["license"].tolist()
+
+            encoded = tokenizer.encode_batch(texts)
+
+            for idx, enc in enumerate(encoded):
+                token_count = len(enc.ids)
+
+                # skip junk (empty or single token)
+                if token_count < 2:
+                    continue
+
+                row = {
+                    "text":        texts[idx],
+                    "token_count": token_count,
+                    "repo":        repos[idx],
+                    "path":        paths[idx],
+                    "license":     licenses[idx],
+                }
+
+                # if this single sample alone exceeds shard size, still include it
+                # — don't lose real data, just let that shard be a bit over
+                if buf["token_count"] + token_count > SHARD_TOKENS and buf["rows"]:
+                    # flush current buffer first
+                    flush_shard(lang, buf["rows"], state)
+                    save_state(state)
+                    buf["rows"] = []
+                    buf["token_count"] = 0
+
+                buf["rows"].append(row)
+                buf["token_count"] += token_count
+
+        state["done"].append(i)
+        save_state(state)
+        print(f"[{i:05d}] ✓ Done")
+
+    # ── Flush any 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 per-language meta ───────────────────────────────────────────────
+    print("\nWriting meta.json per language...")
+    for lang, total_tokens in state["lang_tokens"].items():
+        meta = {
+            "language":     lang,
+            "total_tokens": total_tokens,
+            "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)
+        api.upload_file(
+            path_or_fileobj=str(meta_path),
+            path_in_repo=f"{lang}/meta.json",
+            repo_id=DATASET_REPO,
+            repo_type="dataset",
+            token=HF_TOKEN,
+        )
+        print(f"  {lang}: {total_tokens:,} tokens | {meta['total_shards']} shards")
+
+    print("\n✓ All done!")
+
+# ── Entry point ──────────────────────────────────────────────────────────────
+def main():
+    tokenizer = load_tokenizer()
+    state     = load_state()
+
+    # fire processing in background so Space stays alive
+    t = threading.Thread(target=process, args=(tokenizer, state), daemon=True)
+    t.start()
+
+    # keep the Space running
+    while True:
+        time.sleep(60)
+
+if __name__ == "__main__":
+    main()