File size: 4,312 Bytes
45383d0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """Pull the astro-ph text slice out of ArxivCap via DuckDB (no image download)."""
import argparse
import os
import time
from collections import Counter
from pathlib import Path
import duckdb
SRC = "hf://datasets/MMInstruction/ArxivCap/data/*.parquet"
EXTRACT = """
SELECT
arxiv_id,
title,
abstract,
meta.meta_from_kaggle.categories AS categories,
meta.meta_from_s2.citationCount AS citations,
list_transform(caption_images, x -> x.caption) AS captions,
list_transform(
caption_images,
x -> substr(
array_to_string(
flatten(list_transform(x.cil_pairs, y -> y.image_ocr)), ' '
), 1, 500)
) AS ocr
FROM read_parquet('{src}')
WHERE contains(meta.meta_from_kaggle.categories, 'astro-ph')
{limit}
"""
def year_from_arxiv_id(arxiv_id):
try:
yy = int(arxiv_id.split("/")[1][:2]) if "/" in arxiv_id else int(arxiv_id[:2])
return 1900 + yy if yy > 80 else 2000 + yy
except (ValueError, IndexError):
return None
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--src", default=SRC)
ap.add_argument("--out", default="astro_captions.parquet")
ap.add_argument("--limit", type=int, default=None)
ap.add_argument("--parts-dir", default="census_parts")
args = ap.parse_args()
con = duckdb.connect()
for setting in ["SET enable_progress_bar = true", "SET http_retries = 8",
"SET http_retry_wait_ms = 2000", "SET http_retry_backoff = 2",
"SET threads = 2"]:
try:
con.execute(setting)
except duckdb.Error:
pass
token = os.environ.get("HF_TOKEN", "")
if token:
con.execute(f"CREATE SECRET hf (TYPE HUGGINGFACE, TOKEN '{token}')")
else:
print("WARNING: no HF_TOKEN set, anonymous requests get rate limited fast.")
print("Extracting astro-ph slice (text columns only)...")
if args.limit or not args.src.startswith("hf://"):
limit = f"LIMIT {args.limit}" if args.limit else ""
con.execute(f"COPY ({EXTRACT.format(src=args.src, limit=limit)}) "
f"TO '{args.out}' (FORMAT PARQUET)")
else:
parts = Path(args.parts_dir)
parts.mkdir(exist_ok=True)
chunks = [f"{y:02d}" for y in list(range(91, 100)) + list(range(0, 24))]
for yy in chunks:
part_file = parts / f"astro_{yy}.parquet"
if part_file.exists():
continue
chunk_src = args.src.replace("*", f"arXiv_src_{yy}*")
query = EXTRACT.format(src=chunk_src, limit="")
for attempt in range(5):
try:
con.execute(f"COPY ({query}) TO '{part_file}' (FORMAT PARQUET)")
print(f" {yy}: done")
break
except duckdb.Error as e:
msg = str(e)
if "No files found" in msg:
print(f" {yy}: no files, skipping")
break
if attempt == 4:
raise
wait = 30 * (attempt + 1)
print(f" {yy}: {msg[:80]} -- retrying in {wait}s")
time.sleep(wait)
con.execute(f"COPY (SELECT * FROM read_parquet('{parts}/astro_*.parquet')) "
f"TO '{args.out}' (FORMAT PARQUET)")
papers, figures, usable, fpp, cap_chars = con.execute(f"""
SELECT count(*),
sum(len(captions)),
sum(len(list_filter(captions, c -> length(c) >= 20))),
avg(len(captions)),
round(avg(list_aggregate(list_transform(captions, c -> length(c)), 'avg')), 0)
FROM read_parquet('{args.out}')
""").fetchone()
print(f"\nPapers: {papers:,} Figures: {int(figures):,} "
f"Usable captions: {int(usable):,} Figs/paper: {fpp:.1f} "
f"Mean caption: {int(cap_chars)} chars")
ids = [r[0] for r in con.execute(
f"SELECT arxiv_id FROM read_parquet('{args.out}')").fetchall()]
by_year = Counter(y for y in (year_from_arxiv_id(i) for i in ids) if y)
for y in sorted(by_year):
print(f" {y}: {by_year[y]:,}")
print(f"\nSlice written to {args.out}")
if __name__ == "__main__":
main()
|