Update app.py

app.py

@@ -1,226 +1,10 @@
-import os
-import json
-import time
-import threading
-import io
-import requests
-import pandas as pd
+import shutil
 from pathlib import Path
-from tokenizers import Tokenizer
-from huggingface_hub import HfApi
-import socket
-import threading
-import time
 
-def keep_alive():
-    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)
-    while True:
-        conn, _ = s.accept()
-        conn.send(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nOK")
-        conn.close()
-# ── 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("/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")
+for folder in ["/data/by-language", "/data/codeparrot-raw"]:
+    p = Path(folder)
+    if p.exists():
+        shutil.rmtree(p)
+        print(f"✓ Deleted {folder}")
     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)
-
-# ── 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
-
-        print(f"[{i:05d}] {len(df):,} rows | processing...")
-
-        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()
+        print(f"  Skipped {folder} — not found")