Update app.py

app.py

@@ -1,202 +1,10 @@
-import os
-import json
-import time
-import socket
-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
 
-# ── 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   = 1_000_000  # 1M 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")
+for folder in ["/data/by-language", "/data/progress_state.json"]:
+    p = Path(folder)
+    if p.exists():
+        shutil.rmtree(p) if p.is_dir() else p.unlink()
+        print(f"✓ Deleted {folder}")
     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:05d}.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:05d}/{TOTAL_PARQUETS}] SKIP")
-            continue
-
-        url = PARQUET_URL.format(i=i)
-        print(f"[{i:05d}/{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:05d}] Download error: {e} — skipping")
-            continue
-
-        print(f"[{i:05d}] {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:05d}] ✓ 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)
+        print(f"  Skipped {folder} — not found")