| import json |
| import os |
| import re |
| import shutil |
| import subprocess |
| import sys |
| import time |
| import uuid |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Generator, List, Optional, Tuple |
|
|
|
|
| @dataclass |
| class PipelineConfig: |
| repo_root: Path |
| source_root: Path |
| paper_input: str |
| llm_provider: str |
| llm_model: str |
| llm_model_step4: str |
| model_path: str |
| model_data_dir: str |
| model_class_def: str |
| model_lm: str |
| device: str |
| embedding_model: str |
|
|
|
|
| @dataclass |
| class PipelineResult: |
| job_id: str |
| job_dir: Path |
| paper_dir: Path |
| zip_path: Path |
|
|
|
|
| STEP_LABELS = { |
| 1: "Fetch metadata + LaTeX for input paper", |
| 2: "Add citation markers", |
| 3: "Build usage contexts", |
| 4: "Label citation functions", |
| 5: "Verify USES/EXTENDS", |
| 6: "Extract arXiv paragraphs", |
| 7: "Extract target contributions and refine clusters", |
| } |
|
|
| FULL_STEPS = [1, 2, 3, 4, 5, 6, 7] |
| STOP_PREFIX = "Pipeline stopped:" |
|
|
|
|
| def parse_arxiv_id(paper_input: str) -> str: |
| s = (paper_input or "").strip() |
| if not s: |
| raise ValueError("paper_input is required") |
| if "arxiv.org" in s: |
| m = re.search(r"arxiv\.org/(abs|pdf)/([^/?#]+)", s) |
| if not m: |
| raise ValueError(f"Could not parse arXiv ID from URL: {s}") |
| s = m.group(2) |
| s = s.replace(".pdf", "") |
| s = re.sub(r"v\d+$", "", s) |
| if not re.match(r"^[0-9]{4}\.[0-9]{4,5}$", s): |
| raise ValueError(f"Invalid arXiv ID format: {s}") |
| return s |
|
|
|
|
| def _build_commands( |
| cfg: PipelineConfig, |
| step: int, |
| job_processed_root: Path, |
| paper_id: str, |
| ids_path: Optional[Path], |
| ) -> List[List[str]]: |
| py = sys.executable |
| if step == 1: |
| assert ids_path is not None |
| return [[ |
| py, |
| "src/step_01_fetch/fetch_metadata.py", |
| "--ids", |
| str(ids_path), |
| "--outdir", |
| str(job_processed_root), |
| ]] |
| if step == 2: |
| return [[py, "src/step_02_mark_citations/replace_citation_markers.py", "--root", str(job_processed_root)]] |
| if step == 3: |
| return [[py, "src/step_03_usage_contexts/build_usage_contexts.py", "--root", str(job_processed_root), "--out-name", "usage_contexts.json"]] |
| if step == 4: |
| return [[ |
| py, |
| "src/step_04_label_citations/label_citation_functions.py", |
| "--root", |
| str(job_processed_root), |
| "--model-path", |
| cfg.model_path, |
| "--model-data-dir", |
| cfg.model_data_dir, |
| "--model-class-def", |
| cfg.model_class_def, |
| "--model-lm", |
| cfg.model_lm, |
| "--device", |
| cfg.device, |
| ]] |
| if step == 5: |
| return [[ |
| py, |
| "src/step_05_verify_uses_extends/verify_uses_extends.py", |
| "--root", |
| str(job_processed_root), |
| "--k", |
| "0", |
| "--batch-size", |
| "25", |
| ]] |
| if step == 6: |
| return [[py, "src/step_06_extract_paragraphs/extract_arxiv_paragraphs.py", "--root", str(job_processed_root)]] |
| if step == 7: |
| return [ |
| [py, "src/step_07_extract_and_refine/extract_contributions_from_citations.py", "--root", str(job_processed_root)], |
| [py, "src/step_07_extract_and_refine/refine_and_filter_clusters_llm.py", "--root", str(job_processed_root), "--inplace", "--overwrite"], |
| ] |
| raise ValueError(f"Unknown step: {step}") |
|
|
|
|
| def _write_single_id_file(job_dir: Path, arxiv_id: str) -> Path: |
| ids_path = job_dir / "input_ids.json" |
| payload = [{"id": arxiv_id, "title": "", "id_type": "ArXiv"}] |
| ids_path.write_text(json.dumps(payload, indent=2), encoding="utf-8") |
| return ids_path |
|
|
|
|
| def _write_run_metadata(cfg: PipelineConfig, job_dir: Path, paper_id: str, arxiv_id: str) -> None: |
| payload = { |
| "paper_input": cfg.paper_input, |
| "paper_id": paper_id, |
| "arxiv_id": arxiv_id, |
| "source_root": str(cfg.source_root), |
| "steps": FULL_STEPS + ["annotation"], |
| "llm_provider": cfg.llm_provider, |
| "llm_model": cfg.llm_model, |
| "llm_model_step4": cfg.llm_model_step4, |
| "device": cfg.device, |
| "embedding_model": cfg.embedding_model, |
| "timestamp": int(time.time()), |
| } |
| (job_dir / "run_config.json").write_text(json.dumps(payload, indent=2), encoding="utf-8") |
|
|
|
|
| def _zip_job_dir(job_dir: Path) -> Path: |
| zip_base = job_dir.parent / job_dir.name |
| archive = shutil.make_archive(str(zip_base), "zip", root_dir=str(job_dir)) |
| return Path(archive) |
|
|
|
|
| def _tail_log(path: Path, max_lines: int = 60) -> str: |
| try: |
| lines = path.read_text(encoding="utf-8", errors="ignore").splitlines() |
| except Exception: |
| return "" |
| if not lines: |
| return "" |
| return "\n".join(lines[-max_lines:]) |
|
|
|
|
| def _load_json(path: Path, default=None): |
| try: |
| return json.loads(path.read_text(encoding="utf-8")) |
| except Exception: |
| return default |
|
|
|
|
| def _write_summary_and_zip(job_dir: Path, summary_lines: List[str]) -> Path: |
| (job_dir / "summary.txt").write_text("\n".join(summary_lines), encoding="utf-8") |
| return _zip_job_dir(job_dir) |
|
|
|
|
| def _count_verified_uses_extends(payload: dict) -> int: |
| records = payload.get("confirmed") or payload.get("verified_contexts") or payload.get("contexts") or payload.get("items") or [] |
| if not isinstance(records, list): |
| return 0 |
| accepted = {"USES", "EXTENDS", "Uses", "Extends"} |
| return sum(1 for item in records if isinstance(item, dict) and item.get("label") in accepted) |
|
|
|
|
| def _stop_reason_after_step(step: int, paper_dir: Path) -> str | None: |
| if step == 1: |
| if not paper_dir.exists(): |
| return "metadata could not be fetched for this paper" |
| if not (paper_dir / "processed_main.tex").exists(): |
| return "arXiv source could not be retrieved or converted for this paper" |
| citations = _load_json(paper_dir / "citations_metadata.json", []) |
| if not isinstance(citations, list) or not citations: |
| return "Semantic Scholar returned no citing papers for this target paper" |
|
|
| if step == 3: |
| usage = _load_json(paper_dir / "usage_contexts.json", {}) |
| if not isinstance(usage, dict): |
| return "citation usage contexts could not be built" |
| if int(usage.get("num_contexts") or 0) == 0: |
| return "no citation usage contexts were found" |
|
|
| if step == 4: |
| labels = _load_json(paper_dir / "usage_context_labels.json", {}) |
| contexts = labels.get("labels") if isinstance(labels, dict) else None |
| if not isinstance(contexts, list) or not contexts: |
| return "citation-function labeling produced no labeled contexts" |
|
|
| if step == 5: |
| verified = _load_json(paper_dir / "usage_uses_extends_verified.json", {}) |
| if not isinstance(verified, dict): |
| return "USES/EXTENDS verification did not produce an output file" |
| if _count_verified_uses_extends(verified) == 0: |
| return "no downstream citations were verified as USES or EXTENDS" |
|
|
| if step == 6: |
| paragraphs = _load_json(paper_dir / "usage_citing_paragraphs.json", {}) |
| citing = paragraphs.get("citing_papers") if isinstance(paragraphs, dict) else None |
| if not isinstance(citing, list) or not citing: |
| return "no citing-paper paragraphs could be extracted from arXiv" |
| usable = [ |
| item for item in citing |
| if isinstance(item, dict) |
| and not item.get("error") |
| and (item.get("matched_paragraphs") or item.get("target_citing_paragraphs")) |
| ] |
| if not usable: |
| return "arXiv paragraph extraction returned no usable citing-paper text" |
|
|
| if step == 7: |
| contributions = _load_json(paper_dir / "usage_contributions.json", {}) |
| items = contributions.get("contributions") if isinstance(contributions, dict) else None |
| if not isinstance(items, list) or not items: |
| return "no downstream target-contribution evidence could be extracted" |
| refined = _load_json(paper_dir / "usage_discovery_from_contributions.json", {}) |
| clusters = refined.get("clusters") if isinstance(refined, dict) else None |
| if not isinstance(clusters, list) or not clusters: |
| return "no valid downstream usage clusters survived refinement" |
|
|
| return None |
|
|
|
|
| def run_pipeline(cfg: PipelineConfig, output_root: Path) -> Generator[Tuple[str, Optional[str]], None, PipelineResult]: |
| output_root.mkdir(parents=True, exist_ok=True) |
| job_id = f"job_{int(time.time())}_{uuid.uuid4().hex[:8]}" |
| job_dir = output_root / job_id |
| job_processed_root = job_dir / "processed_papers" |
| job_logs = job_dir / "logs" |
|
|
| job_processed_root.mkdir(parents=True, exist_ok=True) |
| job_logs.mkdir(parents=True, exist_ok=True) |
|
|
| arxiv_id = parse_arxiv_id(cfg.paper_input) |
| paper_id = arxiv_id |
| ids_path = _write_single_id_file(job_dir, arxiv_id) |
| _write_run_metadata(cfg, job_dir, paper_id, arxiv_id) |
|
|
| base_env = os.environ.copy() |
| base_env["LLM_PROVIDER"] = cfg.llm_provider |
| base_env["LLM_MODEL"] = cfg.llm_model |
|
|
| summary_lines: List[str] = [] |
| paper_dir = job_processed_root / paper_id |
|
|
| max_step = 8 |
| for step in FULL_STEPS: |
| label = STEP_LABELS[step] |
| log_file = job_logs / f"step_{step:02d}.log" |
| summary_lines.append(f"[{step}] {label}") |
| yield (f"Step {step}/{max_step}: {label}", None) |
|
|
| env = base_env.copy() |
| if step == 5 and cfg.llm_model_step4: |
| env["LLM_MODEL"] = cfg.llm_model_step4 |
|
|
| with log_file.open("w", encoding="utf-8") as lf: |
| return_code = 0 |
| failed_cmd: List[str] | None = None |
| for cmd in _build_commands(cfg, step, job_processed_root, paper_id, ids_path): |
| lf.write(f"$ {' '.join(cmd)}\n\n") |
| proc = subprocess.Popen( |
| cmd, |
| cwd=str(cfg.repo_root), |
| stdout=subprocess.PIPE, |
| stderr=subprocess.STDOUT, |
| text=True, |
| encoding="utf-8", |
| errors="ignore", |
| env=env, |
| ) |
| assert proc.stdout is not None |
| for line in proc.stdout: |
| lf.write(line) |
| return_code = proc.wait() |
| if return_code != 0: |
| failed_cmd = cmd |
| break |
|
|
| if return_code != 0: |
| summary_lines.append(f"FAILED at step {step}") |
| zip_path = _write_summary_and_zip(job_dir, summary_lines) |
| tail = _tail_log(log_file) |
| if tail: |
| yield ( |
| f"Step {step} failed.\n\nCommand: {' '.join(failed_cmd or [])}\n\nLast log lines:\n{tail}", |
| str(zip_path), |
| ) |
| else: |
| yield (f"Step {step} failed. Command: {' '.join(failed_cmd or [])}", str(zip_path)) |
| return PipelineResult(job_id=job_id, job_dir=job_dir, paper_dir=paper_dir, zip_path=zip_path) |
| else: |
| yield (f"Step {step} complete", None) |
|
|
| if step == 1 and not paper_dir.exists(): |
| summary_lines.append("FAILED: fetch_metadata did not create paper directory") |
| zip_path = _write_summary_and_zip(job_dir, summary_lines) |
| yield (f"Step 1 finished but paper dir missing: {paper_dir}", str(zip_path)) |
| return PipelineResult(job_id=job_id, job_dir=job_dir, paper_dir=paper_dir, zip_path=zip_path) |
|
|
| stop_reason = _stop_reason_after_step(step, paper_dir) |
| if stop_reason: |
| message = f"{STOP_PREFIX} {stop_reason}." |
| summary_lines.append(message) |
| zip_path = _write_summary_and_zip(job_dir, summary_lines) |
| yield (message, str(zip_path)) |
| return PipelineResult(job_id=job_id, job_dir=job_dir, paper_dir=paper_dir, zip_path=zip_path) |
|
|
| summary_lines.append("SUCCESS") |
| zip_path = _write_summary_and_zip(job_dir, summary_lines) |
| yield ("Pipeline completed successfully.", str(zip_path)) |
| return PipelineResult(job_id=job_id, job_dir=job_dir, paper_dir=paper_dir, zip_path=zip_path) |
|
|
