hf_space/runner.py

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)