Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """ | |
| Agentic Pipeline Phase 0: Document Parsing and Indexing. | |
| Pure text processing, no API calls. | |
| 1. Parse paper text via PyMuPDF | |
| 2. Detect sections (numbered, unnumbered, common ML headings) | |
| 3. Chunk at paragraph level (256-512 tokens) respecting section boundaries | |
| 4. Build keyword index for all 30 Croissant fields | |
| 5. Save structured JSON per paper + aggregate summary | |
| """ | |
| import json | |
| import re | |
| from pathlib import Path | |
| from collections import defaultdict | |
| from croissantminer.pdf.reader import extract_text_from_pdf as _canonical_extract_text | |
| from croissantminer.pdf.processor import clean_text as _canonical_clean_text | |
| ROOT = Path(__file__).parent.parent | |
| RAW = ROOT / "data" / "raw" | |
| PROCESSED = ROOT / "data" / "processed" | |
| OUTPUT = ROOT / "data" / "agentic" / "phase0" | |
| OUTPUT.mkdir(parents=True, exist_ok=True) | |
| with open(ROOT / "data" / "paper_links.json") as f: | |
| paper_links = json.load(f) | |
| DUPES = {"CIFAR", "FLORES", "MLS", "MMLU", "MMMU", "MSCOCO", "MathVista", "Visual Genome"} | |
| # ═══════════════════════════════════════════════════════════════ | |
| # KEYWORD DEFINITIONS (all 30 fields) | |
| # ═══════════════════════════════════════════════════════════════ | |
| FIELD_KEYWORDS = { | |
| # General (10) | |
| "name": [ | |
| "dataset", "benchmark", "corpus", "collection", "we present", "we introduce", | |
| "we release", "we propose", "called", "named", "entitled", | |
| ], | |
| "description": [ | |
| "consists of", "contains", "comprising", "composed of", "totaling", | |
| "includes", "encompasses", "covers", "spans", | |
| ], | |
| "url": [ | |
| "github.com", "huggingface", "zenodo", "available at", "download", | |
| "http://", "https://", "repository", "hosted at", "accessible", | |
| ], | |
| "license": [ | |
| "license", "licence", "cc-by", "cc by", "mit license", "apache", | |
| "creative commons", "gpl", "open source", "copyright", "bsd", | |
| "public domain", "proprietary", "non-commercial", | |
| ], | |
| "creator": [ | |
| "created by", "developed by", "constructed by", "curated by", | |
| "compiled by", "authored by", "built by", "prepared by", | |
| ], | |
| "publisher": [ | |
| "published", "proceedings", "conference", "journal", "workshop", | |
| "neurips", "icml", "iclr", "acl", "emnlp", "naacl", "cvpr", "iccv", | |
| "eccv", "aaai", "ijcai", "arxiv", "ieee", "springer", | |
| ], | |
| "datePublished": [ | |
| "released in", "published in", "available since", "introduced in", | |
| "first released", "made public", "released on", | |
| ], | |
| "inLanguage": [ | |
| "english", "multilingual", "monolingual", "bilingual", "language", | |
| "chinese", "french", "german", "spanish", "arabic", "hindi", | |
| "translated", "translation", "cross-lingual", | |
| ], | |
| "citeAs": [ | |
| "cite", "citation", "bibtex", "@article", "@inproceedings", | |
| "please cite", "when using", "reference", | |
| ], | |
| "isLiveDataset": [ | |
| "live", "continuously updated", "maintained", "growing", "evolving", | |
| "static", "frozen", "snapshot", "fixed version", "will not change", | |
| ], | |
| # RAI (20) | |
| "rai:dataCollection": [ | |
| "collect", "gathered", "obtained", "sourced", "crawl", "scrape", | |
| "download", "survey", "interview", "acquisition", "assembled", | |
| "compiled", "harvested", "mined", | |
| ], | |
| "rai:dataCollectionType": [ | |
| "crowdsourc", "web scraping", "web crawl", "manual curation", | |
| "survey", "api", "web api", "experiment", "measurement", | |
| "secondary data", "document analysis", "user-generated", | |
| "self-report", "focus group", "passive collection", | |
| ], | |
| "rai:dataCollectionMissingData": [ | |
| "missing data", "missing value", "incomplete", "absent", "gap", | |
| "unavailable", "n/a", "null", "empty field", "not available", | |
| "data loss", "dropout", "attrition", | |
| ], | |
| "rai:dataCollectionRawData": [ | |
| "raw data", "source data", "original data", "underlying data", | |
| "unprocessed", "original source", "primary source", "obtained from", | |
| ], | |
| "rai:dataCollectionTimeframe": [ | |
| "timeframe", "time frame", "period", "duration", "between.*and", | |
| "from.*to", "start date", "end date", "collected during", | |
| "months", "years", "weeks", "2018", "2019", "2020", "2021", | |
| "2022", "2023", "2024", "2025", | |
| ], | |
| "rai:dataImputationProtocol": [ | |
| "imput", "filled", "replaced missing", "interpolat", "mean substitut", | |
| "median substitut", "knn imputation", "mice", | |
| ], | |
| "rai:dataManipulationProtocol": [ | |
| "augment", "transform", "manipulat", "modification", "synthetic", | |
| "oversample", "undersample", "smote", "rotation", "flip", | |
| "crop", "noise injection", "data augmentation", | |
| ], | |
| "rai:dataPreprocessingProtocol": [ | |
| "preprocess", "cleaning", "filtering", "normaliz", "tokeniz", | |
| "removed", "excluded", "dedup", "deduplication", "lowercas", | |
| "stemm", "lemmatiz", "stopword", "outlier removal", | |
| ], | |
| "rai:dataAnnotationProtocol": [ | |
| "annotat", "label", "tagged", "crowdwork", "crowdsourc", | |
| "mechanical turk", "mturk", "amt", "prolific", "rater", | |
| "judge", "assessor", "gold standard", "ground truth", | |
| "annotation guideline", "codebook", "inter-annotator", | |
| ], | |
| "rai:dataAnnotationPlatform": [ | |
| "mechanical turk", "mturk", "amazon", "scale ai", "surge", | |
| "prolific", "labelbox", "appen", "figure eight", "crowdflower", | |
| "toloka", "platform", "annotation tool", "labeling tool", | |
| ], | |
| "rai:dataAnnotationAnalysis": [ | |
| "inter-annotator", "agreement", "kappa", "fleiss", "cohen", | |
| "krippendorff", "iaa", "reliability", "consensus", "adjudicat", | |
| "disagreement", "conflict resolution", "majority vote", | |
| ], | |
| "rai:annotationsPerItem": [ | |
| "per item", "per example", "per instance", "per sample", | |
| "annotations each", "labels each", "redundan", "multiple annotator", | |
| "three annotator", "two annotator", "five annotator", | |
| ], | |
| "rai:annotatorDemographics": [ | |
| "demographic", "annotator background", "worker", "native speaker", | |
| "qualification", "undergraduate", "graduate", "expert", | |
| "age", "gender", "location", "education", "experience", | |
| ], | |
| "rai:machineAnnotationTools": [ | |
| "automatic", "automated", "model-generated", "gpt", "bert", | |
| "classifier", "detector", "ocr", "ner", "named entity", | |
| "sentiment analysis", "machine label", "auto-label", | |
| "pre-trained", "spacy", "stanza", "nltk", | |
| ], | |
| "rai:dataReleaseMaintenancePlan": [ | |
| "maintenance", "version", "update", "deprecat", "roadmap", | |
| "release plan", "quarterly", "annually", "long-term", | |
| "will be updated", "future version", "maintained by", | |
| ], | |
| "rai:personalSensitiveInformation": [ | |
| "personal", "sensitive", "pii", "personally identifiable", | |
| "identif", "anonymous", "de-identif", "pseudonym", "consent", | |
| "gdpr", "privacy", "confidential", "protected", "hipaa", | |
| "age", "gender", "race", "ethnicity", "religion", "income", | |
| ], | |
| "rai:dataSocialImpact": [ | |
| "social impact", "societal", "broader impact", "harmful", | |
| "benefit", "dual use", "misuse", "ethical implication", | |
| "consequence", "responsibility", "accountability", | |
| "downstream effect", "unintended", | |
| ], | |
| "rai:dataBiases": [ | |
| "bias", "biased", "fairness", "fair", "demographic", "representation", | |
| "skew", "imbalance", "underrepresent", "overrepresent", | |
| "disparity", "equity", "discriminat", "stereotype", | |
| "systematic error", "sampling bias", "selection bias", | |
| ], | |
| "rai:dataLimitations": [ | |
| "limitation", "caveat", "shortcoming", "drawback", "weakness", | |
| "constrain", "not suitable", "does not", "cannot", "restrict", | |
| "scope", "out of scope", "future work", | |
| ], | |
| "rai:dataUseCases": [ | |
| "use case", "intended use", "application", "training", "evaluation", | |
| "benchmark", "fine-tun", "testing", "research purpose", | |
| "not intended for", "suitable for", "designed for", | |
| ], | |
| } | |
| # ═══════════════════════════════════════════════════════════════ | |
| # SECTION DETECTION | |
| # ═══════════════════════════════════════════════════════════════ | |
| # Common ML paper section names (case-insensitive matching) | |
| KNOWN_SECTIONS = { | |
| "abstract", "introduction", "background", "related work", "prior work", | |
| "motivation", "overview", "problem", "formulation", "task", | |
| "method", "methodology", "approach", "system", "model", "architecture", | |
| "pipeline", "framework", "algorithm", | |
| "dataset", "data", "corpus", "benchmark", "collection", "annotation", | |
| "experiment", "experimental", "setup", "evaluation", "result", | |
| "analysis", "ablation", "comparison", | |
| "discussion", "finding", "insight", | |
| "limitation", "future work", "conclusion", "summary", | |
| "ethic", "broader impact", "social impact", "responsible", | |
| "acknowledgment", "reference", "appendix", "supplement", | |
| } | |
| def detect_sections(text): | |
| """Detect section headings and return list of {name, char_start, char_end}.""" | |
| lines = text.split("\n") | |
| sections = [] | |
| char_pos = 0 | |
| for i, line in enumerate(lines): | |
| stripped = line.strip() | |
| line_start = char_pos | |
| char_pos += len(line) + 1 # +1 for newline | |
| if not stripped or len(stripped) > 120: | |
| continue | |
| is_section = False | |
| section_name = stripped | |
| # Pattern 1: Numbered section "1. Introduction" or "1 Introduction" or "3.2 Dataset" | |
| m = re.match(r'^(\d+\.?\d*\.?\s+)(.+)$', stripped) | |
| if m and len(m.group(2)) < 80: | |
| candidate = m.group(2).strip() | |
| # Verify it looks like a heading (not a sentence) | |
| if not candidate.endswith('.') and not candidate.endswith(','): | |
| is_section = True | |
| section_name = candidate | |
| # Pattern 2: ALL CAPS line (common for some formats) | |
| if not is_section and stripped.isupper() and 3 < len(stripped) < 60: | |
| is_section = True | |
| section_name = stripped.title() | |
| # Pattern 3: Known section name — must be a standalone short line (not a sentence) | |
| if not is_section and len(stripped) < 60: | |
| lower = stripped.lower() | |
| clean = re.sub(r'^\d+[\.\d]*\s*', '', lower).strip() | |
| # Must START with a known section name (not just contain it mid-sentence) | |
| for known in KNOWN_SECTIONS: | |
| if clean == known or clean.startswith(known + " ") or clean.startswith(known + ":"): | |
| # Extra check: line should not end with sentence punctuation | |
| if not stripped.endswith('.') and not stripped.endswith(',') and not stripped.endswith(';'): | |
| is_section = True | |
| section_name = stripped | |
| break | |
| # Pattern 4 removed — too aggressive, catches figure captions, equations, etc. | |
| if is_section: | |
| # Clean section name | |
| section_name = re.sub(r'^\d+[\.\d]*\s*', '', section_name).strip() | |
| if section_name: | |
| sections.append({ | |
| "name": section_name, | |
| "char_start": line_start, | |
| }) | |
| # Set char_end for each section | |
| for i in range(len(sections) - 1): | |
| sections[i]["char_end"] = sections[i + 1]["char_start"] | |
| if sections: | |
| sections[-1]["char_end"] = len(text) | |
| # Deduplicate consecutive sections with same name | |
| deduped = [] | |
| for s in sections: | |
| if not deduped or s["name"].lower() != deduped[-1]["name"].lower(): | |
| deduped.append(s) | |
| else: | |
| deduped[-1]["char_end"] = s["char_end"] | |
| return deduped | |
| # ═══════════════════════════════════════════════════════════════ | |
| # CHUNKING | |
| # ═══════════════════════════════════════════════════════════════ | |
| def estimate_tokens(text): | |
| return len(text) // 4 | |
| def chunk_text(text, sections, min_tokens=128, max_tokens=512): | |
| """Split text into chunks respecting section boundaries.""" | |
| if not sections: | |
| sections = [{"name": "Untitled", "char_start": 0, "char_end": len(text)}] | |
| chunks = [] | |
| chunk_idx = 0 | |
| for section in sections: | |
| section_text = text[section["char_start"]:section["char_end"]] | |
| paragraphs = re.split(r'\n\s*\n|\n(?=[A-Z])', section_text) | |
| current_chunk = [] | |
| current_tokens = 0 | |
| current_start = section["char_start"] | |
| for para in paragraphs: | |
| para = para.strip() | |
| if not para: | |
| continue | |
| para_tokens = estimate_tokens(para) | |
| # If adding this paragraph exceeds max, flush | |
| if current_tokens + para_tokens > max_tokens and current_tokens >= min_tokens: | |
| chunk_text_joined = "\n".join(current_chunk) | |
| chunks.append({ | |
| "text": chunk_text_joined, | |
| "section": section["name"], | |
| "chunk_index": chunk_idx, | |
| "token_count": current_tokens, | |
| "char_start": current_start, | |
| "char_end": current_start + len(chunk_text_joined), | |
| }) | |
| chunk_idx += 1 | |
| current_chunk = [] | |
| current_tokens = 0 | |
| current_start = current_start + len(chunk_text_joined) | |
| current_chunk.append(para) | |
| current_tokens += para_tokens | |
| # Flush remaining | |
| if current_chunk: | |
| chunk_text_joined = "\n".join(current_chunk) | |
| chunks.append({ | |
| "text": chunk_text_joined, | |
| "section": section["name"], | |
| "chunk_index": chunk_idx, | |
| "token_count": current_tokens, | |
| "char_start": current_start, | |
| "char_end": current_start + len(chunk_text_joined), | |
| }) | |
| chunk_idx += 1 | |
| return chunks | |
| # ═══════════════════════════════════════════════════════════════ | |
| # KEYWORD INDEX | |
| # ═══════════════════════════════════════════════════════════════ | |
| def build_keyword_index(text, sections): | |
| """Find keyword hits per field, including which sections they appear in.""" | |
| text_lower = text.lower() | |
| index = {} | |
| for field, keywords in FIELD_KEYWORDS.items(): | |
| found_keywords = [] | |
| total_hits = 0 | |
| sections_with_hits = set() | |
| for kw in keywords: | |
| kw_lower = kw.lower() | |
| count = text_lower.count(kw_lower) | |
| if count > 0: | |
| found_keywords.append(kw) | |
| total_hits += count | |
| # Find which sections contain this keyword | |
| for sec in sections: | |
| sec_text = text[sec["char_start"]:sec["char_end"]].lower() | |
| if kw_lower in sec_text: | |
| sections_with_hits.add(sec["name"]) | |
| index[field] = { | |
| "keywords_found": found_keywords, | |
| "total_hits": total_hits, | |
| "sections_found_in": sorted(sections_with_hits), | |
| } | |
| return index | |
| # ═══════════════════════════════════════════════════════════════ | |
| # PAPER TEXT LOADING | |
| # ═══════════════════════════════════════════════════════════════ | |
| def get_paper_text(ds_id): | |
| """Load paper text from PDF.""" | |
| # Try direct name | |
| for name in [f"{ds_id}.pdf"]: | |
| p = RAW / name | |
| if p.exists(): | |
| text = _canonical_clean_text(_canonical_extract_text(str(p))) | |
| return text | |
| # Try arxiv ID from paper_links | |
| url = paper_links.get(ds_id, "") | |
| m = re.search(r'(\d{4}\.\d{4,5})', url) | |
| if m: | |
| arxiv = m.group(1) | |
| for suffix in ["", "v1", "v2", "v3", "v4", "v5"]: | |
| p = RAW / f"{arxiv}{suffix}.pdf" | |
| if p.exists(): | |
| text = _canonical_clean_text(_canonical_extract_text(str(p))) | |
| return text | |
| return None | |
| # ═══════════════════════════════════════════════════════════════ | |
| # MAIN | |
| # ═══════════════════════════════════════════════════════════════ | |
| def process_paper(ds_id): | |
| text = get_paper_text(ds_id) | |
| if text is None: | |
| return None | |
| sections = detect_sections(text) | |
| chunks = chunk_text(text, sections) | |
| keyword_index = build_keyword_index(text, sections) | |
| total_tokens = sum(c["token_count"] for c in chunks) | |
| if total_tokens < 8000: | |
| length_cat = "short" | |
| elif total_tokens < 20000: | |
| length_cat = "medium" | |
| else: | |
| length_cat = "long" | |
| return { | |
| "dataset_id": ds_id, | |
| "chunks": chunks, | |
| "keyword_hits": keyword_index, | |
| "paper_stats": { | |
| "total_tokens": total_tokens, | |
| "total_chunks": len(chunks), | |
| "total_chars": len(text), | |
| "total_sections": len(sections), | |
| "sections_detected": [s["name"] for s in sections], | |
| "paper_length_category": length_cat, | |
| "fields_with_keywords": sum(1 for v in keyword_index.values() if v["keywords_found"]), | |
| }, | |
| } | |
| def main(): | |
| all_ds = sorted( | |
| p.parent.name for p in PROCESSED.glob("*/full_pdf_metadata_result.json") | |
| if p.parent.name not in DUPES | |
| ) | |
| if not all_ds: | |
| # Public release: data/processed/ is not shipped, so take the benchmark papers from the | |
| # dev/test split; their text is read from the PDFs in data/raw/ (scripts/download_papers.py). | |
| split = json.loads((ROOT / "data" / "agentic" / "dev_test_split.json").read_text()) | |
| all_ds = sorted(set(split["dev"]) | set(split["test"])) | |
| print(f"Papers to process: {len(all_ds)}") | |
| # ── Test on 3 ── | |
| test_ids = ["AI4Math_MathVista", "openai_gsm8k", "rajpurkar_squad"] | |
| print(f"\n{'='*70}") | |
| print("VALIDATION: Testing on 3 papers") | |
| print(f"{'='*70}") | |
| for ds_id in test_ids: | |
| result = process_paper(ds_id) | |
| if result is None: | |
| print(f"\n {ds_id}: NO TEXT FOUND") | |
| continue | |
| stats = result["paper_stats"] | |
| print(f"\n {ds_id}:") | |
| print(f" Tokens: {stats['total_tokens']:,} ({stats['paper_length_category']})") | |
| print(f" Chunks: {stats['total_chunks']}") | |
| print(f" Sections: {stats['total_sections']} — {stats['sections_detected'][:8]}{'...' if len(stats['sections_detected']) > 8 else ''}") | |
| print(f" Fields with keywords: {stats['fields_with_keywords']}/30") | |
| # Validate chunks | |
| empty = [c for c in result["chunks"] if not c["text"].strip()] | |
| oversized = [c for c in result["chunks"] if c["token_count"] > 1000] | |
| print(f" Validation: {len(empty)} empty chunks, {len(oversized)} oversized (>1000 tokens)") | |
| # Show top keyword fields | |
| top_kw = sorted(result["keyword_hits"].items(), key=lambda x: -x[1]["total_hits"])[:5] | |
| top_str = ", ".join(str(f.split(":")[-1]) + "(" + str(v["total_hits"]) + ")" for f, v in top_kw) | |
| print(f" Top keyword fields: {top_str}") | |
| # ── Full run ── | |
| print(f"\n{'='*70}") | |
| print(f"Processing all {len(all_ds)} papers...") | |
| print(f"{'='*70}") | |
| success = 0 | |
| failed = 0 | |
| all_stats = [] | |
| field_hit_counts = defaultdict(int) | |
| length_dist = defaultdict(int) | |
| section_fail = [] | |
| for i, ds_id in enumerate(all_ds): | |
| result = process_paper(ds_id) | |
| if result is None: | |
| failed += 1 | |
| continue | |
| out_path = OUTPUT / f"{ds_id}.json" | |
| with open(out_path, "w") as f: | |
| json.dump(result, f, indent=2, ensure_ascii=False) | |
| stats = result["paper_stats"] | |
| all_stats.append(stats) | |
| length_dist[stats["paper_length_category"]] += 1 | |
| for field, hits in result["keyword_hits"].items(): | |
| if hits["keywords_found"]: | |
| field_hit_counts[field] += 1 | |
| if stats["total_sections"] == 0: | |
| section_fail.append(ds_id) | |
| # Validate | |
| empty = [c for c in result["chunks"] if not c["text"].strip()] | |
| oversized = [c for c in result["chunks"] if c["token_count"] > 1000] | |
| if empty or oversized: | |
| print(f" [{i+1}] {ds_id}: {len(empty)} empty, {len(oversized)} oversized") | |
| success += 1 | |
| if (i + 1) % 25 == 0: | |
| print(f" [{i+1}/{len(all_ds)}] {success} OK, {failed} failed") | |
| print(f" [{len(all_ds)}/{len(all_ds)}] {success} OK, {failed} failed") | |
| # ── Summary ── | |
| total_chunks = sum(s["total_chunks"] for s in all_stats) | |
| total_tokens = sum(s["total_tokens"] for s in all_stats) | |
| avg_sections = sum(s["total_sections"] for s in all_stats) / max(len(all_stats), 1) | |
| print(f"\n{'='*70}") | |
| print("PHASE 0 SUMMARY") | |
| print(f"{'='*70}") | |
| print(f" Papers processed: {success}/{len(all_ds)}") | |
| print(f" Failed (no PDF): {failed}") | |
| print(f" Total chunks: {total_chunks:,}") | |
| print(f" Total tokens: {total_tokens:,}") | |
| print(f" Avg chunks/paper: {total_chunks/max(success,1):.0f}") | |
| print(f" Avg tokens/paper: {total_tokens/max(success,1):,.0f}") | |
| print(f" Avg sections/paper: {avg_sections:.1f}") | |
| print(f" Section detection failures: {len(section_fail)}") | |
| if section_fail: | |
| print(f" Papers with 0 sections: {section_fail[:5]}{'...' if len(section_fail) > 5 else ''}") | |
| print(f" Paper length distribution: {dict(length_dist)}") | |
| print(f"\n Per-field keyword hit rate (% of papers with keywords):") | |
| for field in sorted(FIELD_KEYWORDS.keys()): | |
| rate = field_hit_counts[field] / max(success, 1) * 100 | |
| short = field.split(":")[-1] if ":" in field else field | |
| bar = "#" * int(rate / 5) | |
| print(f" {short:<30} {rate:>5.1f}% {bar}") | |
| # Save summary | |
| summary = { | |
| "phase": 0, | |
| "description": "Document Parsing and Indexing", | |
| "papers_processed": success, | |
| "papers_failed": failed, | |
| "total_chunks": total_chunks, | |
| "total_tokens": total_tokens, | |
| "avg_chunks_per_paper": round(total_chunks / max(success, 1), 1), | |
| "avg_tokens_per_paper": round(total_tokens / max(success, 1)), | |
| "avg_sections_per_paper": round(avg_sections, 1), | |
| "section_detection_failures": len(section_fail), | |
| "paper_length_distribution": dict(length_dist), | |
| "field_keyword_hit_rates": { | |
| field: round(field_hit_counts[field] / max(success, 1) * 100, 1) | |
| for field in sorted(FIELD_KEYWORDS.keys()) | |
| }, | |
| } | |
| summary_path = ROOT / "results" / "agentic_phase0_summary.json" | |
| with open(summary_path, "w") as f: | |
| json.dump(summary, f, indent=2) | |
| print(f"\n Summary saved: {summary_path}") | |
| # Verify output files | |
| json_files = list(OUTPUT.glob("*.json")) | |
| print(f" Output files: {len(json_files)}/103") | |
| if __name__ == "__main__": | |
| main() | |