code/final_unique.py

"""Unique counts after filtering: canonical skills x near-dup clusters (train_ok / non-CJK). Downloads each canonical part,
keeps 3 columns, deletes it. Writes logs/final_unique.json and stats/final_unique.json in the bucket."""
import json, os, subprocess, sys
from pathlib import Path
import pyarrow as pa, pyarrow.parquet as pq
ROOT = Path(__file__).resolve().parent.parent
os.environ.setdefault("HOME", "/root")
from dotenv import load_dotenv; load_dotenv(ROOT / ".env")
from huggingface_hub import HfApi
api = HfApi(); tmp = ROOT / "work/final_unique"; tmp.mkdir(exist_ok=True); cols = []
for h in "0123456789abcdef":
    local = tmp / f"skills-{h}.parquet"
    code = ("from huggingface_hub import HfApi; from dotenv import load_dotenv; load_dotenv('/root/skills-db/.env');"
            f"HfApi().download_bucket_files('Mercity/SkillsStorage', [('processed/canonical/skills/skills-{h}.parquet', '{local}')])")
    subprocess.run([sys.executable, "-c", code], check=True)
    cols.append(pq.read_table(local, columns=["sha", "train_ok", "flag_cjk"])); local.unlink()
    print("read", h, flush=True)
t = pa.concat_tables(cols); cl = pq.read_table(ROOT / "work/minhash/clusters.parquet", columns=["sha", "cluster_id"])
j = t.join(cl, keys="sha", join_type="left outer")
import pyarrow.compute as pc
def uniq(mask): return len(pc.unique(pc.filter(j["cluster_id"], mask)))
out = {"canonical_skills": j.num_rows, "clusters_all": len(pc.unique(j["cluster_id"])),
       "train_ok_skills": int(pc.sum(j["train_ok"]).as_py()), "train_ok_clusters": uniq(j["train_ok"]),
       "non_cjk_skills": int(pc.sum(pc.invert(j["flag_cjk"])).as_py()), "non_cjk_clusters": uniq(pc.invert(j["flag_cjk"]))}
(ROOT / "logs/final_unique.json").write_text(json.dumps(out, indent=1)); print(json.dumps(out), flush=True)
api.batch_bucket_files("Mercity/SkillsStorage", add=[(str(ROOT / "logs/final_unique.json"), "stats/final_unique.json")])
print("done", flush=True)