Spaces:
Running on Zero
Running on Zero
File size: 5,664 Bytes
9936912 | 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 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | #!/usr/bin/env python3
"""List or download discovered arXiv PDFs with resume and provenance records."""
from __future__ import annotations
import argparse
import hashlib
import json
import re
import time
import urllib.request
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
PROJECT_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_INPUT = PROJECT_ROOT / "data" / "sources" / "discovered_arxiv.jsonl"
DEFAULT_OUTPUT_DIR = PROJECT_ROOT / "data" / "raw" / "arxiv"
def file_sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as stream:
for block in iter(lambda: stream.read(1024 * 1024), b""):
digest.update(block)
return digest.hexdigest()
def safe_filename(arxiv_id: str) -> str:
return re.sub(r"[^A-Za-z0-9._-]+", "_", arxiv_id) + ".pdf"
def load_records(path: Path) -> list[dict]:
return [json.loads(line) for line in path.open(encoding="utf-8")]
def select_records(
records: list[dict],
excluded_domains: set[str],
max_papers: int | None,
per_domain: int | None,
) -> list[dict]:
selected = [
record
for record in records
if set(record["matched_domains"]) - excluded_domains
]
selected.sort(key=lambda row: (row["matched_domains"], row["title"]))
if per_domain is not None:
balanced = {}
domains = sorted(
{domain for record in selected for domain in record["matched_domains"]}
- excluded_domains
)
for domain in domains:
matches = [record for record in selected if domain in record["matched_domains"]]
for record in matches[:per_domain]:
balanced[record["arxiv_id"]] = record
selected = sorted(balanced.values(), key=lambda row: (row["matched_domains"], row["title"]))
return selected[:max_papers] if max_papers is not None else selected
def download_pdf(record: dict, output_dir: Path) -> dict:
filename = safe_filename(record["arxiv_id"])
destination = output_dir / filename
partial = destination.with_suffix(".pdf.part")
if not destination.exists():
request = urllib.request.Request(
record["pdf_url"], headers={"User-Agent": "controlai-literature-fetcher/0.1"}
)
try:
with urllib.request.urlopen(request, timeout=120) as response, partial.open("wb") as stream:
while block := response.read(1024 * 1024):
stream.write(block)
if partial.read_bytes()[:5] != b"%PDF-":
raise ValueError("downloaded content is not a PDF")
partial.replace(destination)
except Exception:
partial.unlink(missing_ok=True)
raise
return {
"source_id": f"arxiv_{record['arxiv_id']}",
"arxiv_id": record["arxiv_id"],
"title": record["title"],
"authors": record.get("authors", []),
"filename": filename,
"url": record["pdf_url"],
"pdf_url": record["pdf_url"],
"corpus_tier": "research_paper",
"coverage": record["matched_domains"],
"matched_domains": record["matched_domains"],
"license": record["license_url"],
"license_url": record["license_url"],
"bytes": destination.stat().st_size,
"sha256": file_sha256(destination),
"downloaded_at": datetime.now(timezone.utc).isoformat(),
}
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--input", type=Path, default=DEFAULT_INPUT)
parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT_DIR)
parser.add_argument("--exclude-domain", action="append", default=["foundations"])
parser.add_argument("--max-papers", type=int, default=None)
parser.add_argument("--per-domain", type=int, default=None)
parser.add_argument("--delay-seconds", type=float, default=3.2)
parser.add_argument("--download", action="store_true")
args = parser.parse_args()
records = select_records(
load_records(args.input), set(args.exclude_domain), args.max_papers, args.per_domain
)
domain_counts = Counter(domain for row in records for domain in row["matched_domains"])
print(f"Selected papers: {len(records)}")
print("Domains:", json.dumps(dict(sorted(domain_counts.items())), indent=2))
print(f"Output: {args.output_dir}")
if not args.download:
print("Dry run only. Add --download to fetch PDFs.")
return
args.output_dir.mkdir(parents=True, exist_ok=True)
lock_path = args.output_dir / "source_lock.json"
existing = {}
if lock_path.exists():
existing = {row["arxiv_id"]: row for row in json.loads(lock_path.read_text(encoding="utf-8"))}
failures = []
for index, record in enumerate(records, start=1):
try:
existing[record["arxiv_id"]] = download_pdf(record, args.output_dir)
status = "ok"
except Exception as error:
failures.append({"arxiv_id": record["arxiv_id"], "error": str(error)})
status = f"error: {error}"
lock_path.write_text(
json.dumps(sorted(existing.values(), key=lambda row: row["arxiv_id"]), indent=2) + "\n",
encoding="utf-8",
)
print(f"[{index}/{len(records)}] {record['arxiv_id']}: {status}")
if index < len(records):
time.sleep(args.delay_seconds)
print(f"Downloaded/verified: {len(existing)}")
print(f"Failures: {len(failures)}")
if failures:
print(json.dumps(failures, indent=2))
if __name__ == "__main__":
main()
|