| |
| """SERIAL harvester: fetch -> extract to compact jsonl.gz -> DELETE raw. Disk stays flat. |
| One source at a time. Never runs concurrently with itself.""" |
| import os,sys,json,gzip,glob,time,shutil,subprocess,zipfile,io,csv,re |
| os.environ.setdefault("HF_HOME","/workspace/.hf_home") |
| from huggingface_hub import snapshot_download |
| TOK=os.environ.get("HF_TOKEN") |
| STAGE="/workspace/aditi/_stage"; HARV="/workspace/aditi/harvest" |
| LEDGER="/workspace/aditi/ghdata/_reports/harvest_ledger.json" |
| os.makedirs(HARV,exist_ok=True) |
| MAX_ROWS=int(os.environ.get("MAX_ROWS","1500000")) |
| MIN_FREE_GB=25 |
| TEXTCOLS=re.compile(r"(text|sentence|content|body|claim|question|answer|comment|title|" |
| r"headline|tweet|post|review|source|context|doc|output|response|" |
| r"summary|prompt|instruction|input|target|translation|word|token|caption)",re.I) |
|
|
| def free_gb(): |
| s=os.statvfs("/"); return s.f_bavail*s.f_frsize/1e9 |
|
|
| def rows_from_any(path): |
| """Yield dicts from parquet/csv/tsv/json/jsonl/txt inside a directory tree.""" |
| import pandas as pd |
| exts=("**/*.parquet","**/*.csv","**/*.tsv","**/*.jsonl","**/*.json","**/*.txt") |
| files=[] |
| for e in exts: files+=glob.glob(os.path.join(path,e),recursive=True) |
| files=[f for f in files if ".cache" not in f and os.path.getsize(f)>0] |
| files.sort(key=os.path.getsize) |
| for f in files[:2000]: |
| try: |
| low=f.lower() |
| if low.endswith(".parquet"): |
| d=pd.read_parquet(f) |
| for r in d.head(MAX_ROWS).to_dict("records"): yield r |
| elif low.endswith((".csv",".tsv")): |
| sep="\t" if low.endswith(".tsv") else "," |
| for chunk in pd.read_csv(f,sep=sep,encoding="utf-8-sig",on_bad_lines="skip", |
| chunksize=50000,dtype=str,engine="python"): |
| for r in chunk.to_dict("records"): yield r |
| elif low.endswith(".jsonl"): |
| with open(f,encoding="utf-8",errors="replace") as fh: |
| for i,l in enumerate(fh): |
| if i>=MAX_ROWS: break |
| try: yield json.loads(l) |
| except Exception: pass |
| elif low.endswith(".json"): |
| with open(f,encoding="utf-8",errors="replace") as fh: |
| try: d=json.load(fh) |
| except Exception: continue |
| if isinstance(d,list): |
| for r in d[:MAX_ROWS]: |
| if isinstance(r,dict): yield r |
| elif low.endswith(".txt") and os.path.getsize(f)<400_000_000: |
| with open(f,encoding="utf-8",errors="replace") as fh: |
| for i,l in enumerate(fh): |
| if i>=MAX_ROWS: break |
| l=l.strip() |
| if len(l)>20: yield {"text":l} |
| except Exception: |
| continue |
|
|
| def compact(r): |
| """Keep only text-bearing fields + short labels. Drops blobs/images.""" |
| o={} |
| for k,v in (r.items() if isinstance(r,dict) else []): |
| if v is None: continue |
| if isinstance(v,(bytes,bytearray)): continue |
| s=v if isinstance(v,str) else (json.dumps(v,ensure_ascii=False,default=str) |
| if isinstance(v,(list,dict)) else str(v)) |
| if len(s)>6000: s=s[:6000] |
| if TEXTCOLS.search(str(k)) or len(s)<=64: |
| o[str(k)]=s |
| return o if any(len(v)>15 for v in o.values()) else None |
|
|
| def extract(name, srcdir, meta): |
| out=os.path.join(HARV,f"{name}.jsonl.gz"); n=0 |
| with gzip.open(out,"wt",encoding="utf-8") as w: |
| w.write(json.dumps({"_meta":meta})+"\n") |
| for r in rows_from_any(srcdir): |
| c=compact(r) |
| if not c: continue |
| w.write(json.dumps(c,ensure_ascii=False)+"\n"); n+=1 |
| if n>=MAX_ROWS: break |
| return out,n,os.path.getsize(out) |
|
|
| def do_hf(repo,name,pats=None,lic=""): |
| d=os.path.join(STAGE,name) |
| snapshot_download(repo_id=repo,repo_type="dataset",local_dir=d, |
| allow_patterns=pats,max_workers=8,token=TOK) |
| return d,{"source":repo,"kind":"hf","licence":lic} |
| def do_git(url,name,lic=""): |
| d=os.path.join(STAGE,name) |
| subprocess.run(["git","clone","--depth","1","--quiet",url,d], |
| capture_output=True,timeout=1200,check=True) |
| return d,{"source":url,"kind":"git","licence":lic} |
|
|
| def run(queue): |
| led=[] |
| if os.path.exists(LEDGER): |
| try: led=json.load(open(LEDGER))["items"] |
| except Exception: led=[] |
| done={i["name"] for i in led if i.get("status")=="harvested"} |
| for kind,ref,name,pats,lic in queue: |
| if name in done: |
| print(f"skip {name} (already harvested)",flush=True); continue |
| if free_gb()<MIN_FREE_GB: |
| print(f"ABORT: only {free_gb():.0f}G free",flush=True); break |
| t0=time.time(); shutil.rmtree(STAGE,ignore_errors=True); os.makedirs(STAGE,exist_ok=True) |
| try: |
| d,meta=(do_hf(ref,name,pats,lic) if kind=="hf" else do_git(ref,name,lic)) |
| raw=sum(os.path.getsize(os.path.join(dp,f)) for dp,_,fs in os.walk(d) for f in fs) |
| out,n,sz=extract(name,d,meta) |
| shutil.rmtree(STAGE,ignore_errors=True) |
| led.append({"name":name,"ref":ref,"kind":kind,"licence":lic,"status":"harvested", |
| "rows":n,"raw_mb":round(raw/1e6),"kept_mb":round(sz/1e6,1), |
| "secs":round(time.time()-t0)}) |
| print(f"OK {name:26s} rows={n:>7} raw={raw/1e6:7.0f}MB kept={sz/1e6:6.1f}MB " |
| f"free={free_gb():.0f}G",flush=True) |
| except Exception as e: |
| shutil.rmtree(STAGE,ignore_errors=True) |
| led.append({"name":name,"ref":ref,"kind":kind,"status":"failed", |
| "error":f"{type(e).__name__}: {str(e)[:160]}"}) |
| print(f"FAIL {name:26s} {type(e).__name__}: {str(e)[:110]}",flush=True) |
| json.dump({"generated":time.strftime("%Y-%m-%dT%H:%M:%SZ",time.gmtime()), |
| "items":led},open(LEDGER,"w"),indent=2) |
| ok=[i for i in led if i.get("status")=="harvested"] |
| print(f"\nHARVEST harvested={len(ok)} rows={sum(i.get('rows',0) for i in ok)} " |
| f"kept={sum(i.get('kept_mb',0) for i in ok)/1024:.2f}GB free={free_gb():.0f}G",flush=True) |
| if __name__=="__main__": |
| from harvest_queue import QUEUE |
| run(QUEUE) |
|
|
