Upload plugins/mlintern/skills/ml-intern-harness/scripts/papers.py with huggingface_hub

plugins/mlintern/skills/ml-intern-harness/scripts/papers.py

@@ -0,0 +1,476 @@
+#!/usr/bin/env python3
+"""Paper research helper for ML Intern Codex.
+
+Emulates the useful parts of upstream ml-intern's hf_papers tool with public
+HTTP APIs: Hugging Face Papers, arXiv/ar5iv HTML, and Semantic Scholar.
+"""
+
+from __future__ import annotations
+
+import argparse
+import html.parser
+import json
+import os
+import re
+import sys
+import urllib.error
+import urllib.parse
+import urllib.request
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any
+
+
+HF_API = "https://huggingface.co/api"
+ARXIV_HTML = "https://arxiv.org/html"
+AR5IV_HTML = "https://ar5iv.labs.arxiv.org/html"
+S2_API = "https://api.semanticscholar.org"
+MAX_SECTION_TEXT_LEN = 8000
+
+
+def request_json(url: str, params: dict[str, Any] | None = None, method: str = "GET", body: dict[str, Any] | None = None) -> Any:
+    if params:
+        url = f"{url}?{urllib.parse.urlencode({k: v for k, v in params.items() if v is not None})}"
+    data = None
+    headers = {"User-Agent": "ml-intern-codex/0.1"}
+    if body is not None:
+        data = json.dumps(body).encode("utf-8")
+        headers["Content-Type"] = "application/json"
+    s2_key = os.environ.get("S2_API_KEY")
+    if s2_key and url.startswith(S2_API):
+        headers["x-api-key"] = s2_key
+    request = urllib.request.Request(url, data=data, headers=headers, method=method)
+    try:
+        with urllib.request.urlopen(request, timeout=30) as response:
+            return json.loads(response.read().decode("utf-8"))
+    except urllib.error.HTTPError as exc:
+        text = exc.read().decode("utf-8", errors="replace")
+        raise RuntimeError(f"{url} returned HTTP {exc.code}: {text[:500]}") from exc
+
+
+def request_text(url: str) -> str:
+    request = urllib.request.Request(url, headers={"User-Agent": "ml-intern-codex/0.1"})
+    with urllib.request.urlopen(request, timeout=30) as response:
+        return response.read().decode("utf-8", errors="replace")
+
+
+def arxiv_s2_id(arxiv_id: str) -> str:
+    return f"ARXIV:{arxiv_id}"
+
+
+def truncate(text: str, limit: int) -> str:
+    text = re.sub(r"\s+", " ", text).strip()
+    return text if len(text) <= limit else text[:limit].rstrip() + "..."
+
+
+def paper_arxiv_id(paper: dict[str, Any]) -> str:
+    external = paper.get("externalIds") or paper.get("external_ids") or {}
+    return external.get("ArXiv") or paper.get("arxiv_id") or paper.get("id", "")
+
+
+def format_hf_paper(paper: dict[str, Any], idx: int) -> str:
+    nested = paper.get("paper") if isinstance(paper.get("paper"), dict) else paper
+    title = nested.get("title") or paper.get("title") or "(untitled)"
+    arxiv_id = nested.get("id") or nested.get("arxivId") or paper.get("id") or ""
+    summary = nested.get("summary") or nested.get("abstract") or ""
+    lines = [f"### {idx}. {title}"]
+    if arxiv_id:
+        lines.append(f"arxiv_id: {arxiv_id}")
+        lines.append(f"https://arxiv.org/abs/{arxiv_id}")
+    if nested.get("publishedAt"):
+        lines.append(f"Published: {nested['publishedAt']}")
+    if nested.get("githubUrl"):
+        lines.append(f"GitHub: {nested['githubUrl']}")
+    if summary:
+        lines.append(truncate(summary, 500))
+    return "\n".join(lines)
+
+
+def format_s2_paper(paper: dict[str, Any], idx: int) -> str:
+    title = paper.get("title") or "(untitled)"
+    year = paper.get("year") or "?"
+    cites = paper.get("citationCount", 0)
+    venue = paper.get("venue") or ""
+    arxiv_id = paper_arxiv_id(paper)
+    tldr = (paper.get("tldr") or {}).get("text", "")
+    parts = [f"Year: {year}", f"Citations: {cites}"]
+    if venue:
+        parts.append(f"Venue: {venue}")
+    if arxiv_id:
+        parts.append(f"arxiv_id: {arxiv_id}")
+    lines = [f"### {idx}. {title}", " | ".join(parts)]
+    if arxiv_id:
+        lines.append(f"https://arxiv.org/abs/{arxiv_id}")
+    if tldr:
+        lines.append(f"TL;DR: {tldr}")
+    return "\n".join(lines)
+
+
+class ArxivHTMLParser(html.parser.HTMLParser):
+    def __init__(self) -> None:
+        super().__init__()
+        self.capture_title = False
+        self.capture_abstract = False
+        self.capture_heading = False
+        self.capture_paragraph = False
+        self.title_parts: list[str] = []
+        self.abstract_parts: list[str] = []
+        self.sections: list[dict[str, Any]] = []
+        self.current_heading: list[str] = []
+        self.current_paragraph: list[str] = []
+        self.current_section: dict[str, Any] | None = None
+
+    def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
+        classes = dict(attrs).get("class", "") or ""
+        if tag == "h1" and "ltx_title" in classes:
+            self.capture_title = True
+        elif tag == "div" and "ltx_abstract" in classes:
+            self.capture_abstract = True
+        elif tag in {"h2", "h3"} and "ltx_title" in classes:
+            self.capture_heading = True
+            self.current_heading = []
+        elif tag == "p":
+            self.capture_paragraph = True
+            self.current_paragraph = []
+
+    def handle_endtag(self, tag: str) -> None:
+        if tag == "h1" and self.capture_title:
+            self.capture_title = False
+        elif tag == "div" and self.capture_abstract:
+            self.capture_abstract = False
+        elif tag in {"h2", "h3"} and self.capture_heading:
+            heading = truncate(" ".join(self.current_heading), 500)
+            section_id = ""
+            match = re.match(r"^([A-Z]?\d+(?:\.\d+)*)\s", heading)
+            if match:
+                section_id = match.group(1)
+            self.current_section = {"id": section_id, "title": heading, "text": ""}
+            self.sections.append(self.current_section)
+            self.capture_heading = False
+        elif tag == "p" and self.capture_paragraph:
+            paragraph = truncate(" ".join(self.current_paragraph), 4000)
+            if paragraph:
+                if self.capture_abstract:
+                    self.abstract_parts.append(paragraph)
+                elif self.current_section is not None:
+                    existing = self.current_section["text"]
+                    self.current_section["text"] = (existing + "\n\n" + paragraph).strip()
+            self.capture_paragraph = False
+
+    def handle_data(self, data: str) -> None:
+        text = data.strip()
+        if not text:
+            return
+        if self.capture_title:
+            self.title_parts.append(text.removeprefix("Title:"))
+        if self.capture_heading:
+            self.current_heading.append(text)
+        if self.capture_paragraph:
+            self.current_paragraph.append(text)
+
+
+def parse_arxiv_html(html_text: str) -> dict[str, Any]:
+    parser = ArxivHTMLParser()
+    parser.feed(html_text)
+    return {
+        "title": truncate(" ".join(parser.title_parts), 500),
+        "abstract": truncate(" ".join(parser.abstract_parts), 2000),
+        "sections": parser.sections,
+    }
+
+
+def op_trending(args: argparse.Namespace) -> str:
+    params: dict[str, Any] = {"limit": args.limit * 3 if args.query else args.limit}
+    if args.date:
+        params["date"] = args.date
+    papers = request_json(f"{HF_API}/daily_papers", params)
+    if args.query:
+        needle = args.query.lower()
+        papers = [
+            paper
+            for paper in papers
+            if needle in json.dumps(paper, ensure_ascii=False).lower()
+        ]
+    lines = ["# Trending Papers"]
+    for idx, paper in enumerate(papers[: args.limit], 1):
+        lines.append(format_hf_paper(paper, idx))
+        lines.append("")
+    return "\n".join(lines)
+
+
+def op_search(args: argparse.Namespace) -> str:
+    if not args.query:
+        raise SystemExit("search requires --query")
+    use_s2 = any([args.date_from, args.date_to, args.categories, args.min_citations, args.sort_by != "relevance"])
+    if use_s2:
+        params: dict[str, Any] = {
+            "query": args.query,
+            "limit": args.limit,
+            "fields": "title,externalIds,year,citationCount,tldr,venue,publicationDate",
+        }
+        if args.date_from or args.date_to:
+            params["publicationDateOrYear"] = f"{args.date_from or ''}:{args.date_to or ''}"
+        if args.categories:
+            params["fieldsOfStudy"] = args.categories
+        if args.min_citations:
+            params["minCitationCount"] = str(args.min_citations)
+        if args.sort_by != "relevance":
+            params["sort"] = f"{args.sort_by}:desc"
+        data = request_json(f"{S2_API}/graph/v1/paper/search/bulk", params)
+        papers = data.get("data", [])
+        lines = [f"# Papers matching '{args.query}' (Semantic Scholar)"]
+        for idx, paper in enumerate(papers[: args.limit], 1):
+            lines.append(format_s2_paper(paper, idx))
+            lines.append("")
+        return "\n".join(lines)
+    papers = request_json(f"{HF_API}/papers/search", {"q": args.query, "limit": args.limit})
+    lines = [f"# Papers matching '{args.query}' (Hugging Face Papers)"]
+    for idx, paper in enumerate(papers[: args.limit], 1):
+        lines.append(format_hf_paper(paper, idx))
+        lines.append("")
+    return "\n".join(lines)
+
+
+def op_paper_details(args: argparse.Namespace) -> str:
+    if not args.arxiv_id:
+        raise SystemExit("paper_details requires --arxiv-id")
+    paper = request_json(f"{HF_API}/papers/{args.arxiv_id}")
+    lines = [f"# {paper.get('title', args.arxiv_id)}", f"https://huggingface.co/papers/{args.arxiv_id}", f"https://arxiv.org/abs/{args.arxiv_id}"]
+    for key in ("publishedAt", "submittedOnDailyAt", "githubUrl"):
+        if paper.get(key):
+            lines.append(f"{key}: {paper[key]}")
+    if paper.get("summary"):
+        lines.append("")
+        lines.append("## Abstract")
+        lines.append(paper["summary"])
+    if paper.get("ai_summary"):
+        lines.append("")
+        lines.append("## AI Summary")
+        lines.append(str(paper["ai_summary"]))
+    return "\n".join(lines)
+
+
+def op_read_paper(args: argparse.Namespace) -> str:
+    if not args.arxiv_id:
+        raise SystemExit("read_paper requires --arxiv-id")
+    parsed = None
+    for base in (ARXIV_HTML, AR5IV_HTML):
+        try:
+            parsed = parse_arxiv_html(request_text(f"{base}/{args.arxiv_id}"))
+            if parsed["sections"]:
+                break
+        except Exception:
+            parsed = None
+    if not parsed or not parsed["sections"]:
+        return op_paper_details(args) + f"\n\nHTML sections unavailable. PDF: https://arxiv.org/pdf/{args.arxiv_id}"
+    if not args.section:
+        lines = [f"# {parsed['title'] or args.arxiv_id}", f"https://arxiv.org/abs/{args.arxiv_id}", "", "## Abstract", parsed["abstract"], "", "## Sections"]
+        for section in parsed["sections"]:
+            preview = truncate(section.get("text", ""), 280)
+            lines.append(f"- {section['title']}: {preview}")
+        return "\n".join(lines)
+    wanted = args.section.lower()
+    selected = None
+    for section in parsed["sections"]:
+        if section["id"].lower() == wanted or wanted in section["title"].lower():
+            selected = section
+            break
+    if not selected:
+        available = "\n".join(f"- {section['title']}" for section in parsed["sections"])
+        raise SystemExit(f"section not found. Available sections:\n{available}")
+    return "\n".join([
+        f"# {selected['title']}",
+        f"https://arxiv.org/abs/{args.arxiv_id}",
+        "",
+        truncate(selected.get("text", ""), MAX_SECTION_TEXT_LEN),
+    ])
+
+
+def format_citation(entry: dict[str, Any]) -> str:
+    paper = entry.get("citingPaper") or entry.get("citedPaper") or {}
+    title = paper.get("title") or "(untitled)"
+    year = paper.get("year") or "?"
+    cites = paper.get("citationCount", 0)
+    arxiv_id = paper_arxiv_id(paper)
+    line = f"- {title} ({year}, {cites} cites)"
+    if arxiv_id:
+        line += f" arxiv:{arxiv_id}"
+    if entry.get("isInfluential"):
+        line += " [influential]"
+    contexts = entry.get("contexts") or []
+    if contexts:
+        line += f"\n  > {truncate(contexts[0], 220)}"
+    return line
+
+
+def op_citation_graph(args: argparse.Namespace) -> str:
+    if not args.arxiv_id:
+        raise SystemExit("citation_graph requires --arxiv-id")
+    fields = "title,externalIds,year,citationCount,influentialCitationCount,contexts,intents,isInfluential"
+    params = {"fields": fields, "limit": args.limit}
+    paths: dict[str, str] = {}
+    if args.direction in {"references", "both"}:
+        paths["References"] = f"/graph/v1/paper/{arxiv_s2_id(args.arxiv_id)}/references"
+    if args.direction in {"citations", "both"}:
+        paths["Citations"] = f"/graph/v1/paper/{arxiv_s2_id(args.arxiv_id)}/citations"
+    lines = [f"# Citation Graph for {args.arxiv_id}", f"https://arxiv.org/abs/{args.arxiv_id}"]
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        futures = {pool.submit(request_json, f"{S2_API}{path}", params): name for name, path in paths.items()}
+        for future in as_completed(futures):
+            name = futures[future]
+            lines.append("")
+            lines.append(f"## {name}")
+            try:
+                data = future.result()
+                for entry in data.get("data", []):
+                    lines.append(format_citation(entry))
+            except Exception as exc:
+                lines.append(f"Error: {exc}")
+    return "\n".join(lines)
+
+
+def op_resources(args: argparse.Namespace) -> str:
+    if not args.arxiv_id:
+        raise SystemExit(f"{args.operation} requires --arxiv-id")
+    sort = {"downloads": "downloads", "likes": "likes", "trending": "trendingScore"}[args.sort]
+    calls: dict[str, tuple[str, dict[str, Any]]] = {}
+    if args.operation in {"find_datasets", "find_all_resources"}:
+        calls["Datasets"] = (f"{HF_API}/datasets", {"filter": f"arxiv:{args.arxiv_id}", "limit": args.limit, "sort": sort, "direction": -1})
+    if args.operation in {"find_models", "find_all_resources"}:
+        calls["Models"] = (f"{HF_API}/models", {"filter": f"arxiv:{args.arxiv_id}", "limit": args.limit, "sort": sort, "direction": -1})
+    if args.operation in {"find_collections", "find_all_resources"}:
+        calls["Collections"] = (f"{HF_API}/collections", {"paper": args.arxiv_id})
+    lines = [f"# Resources linked to paper {args.arxiv_id}", f"https://huggingface.co/papers/{args.arxiv_id}"]
+    with ThreadPoolExecutor(max_workers=3) as pool:
+        futures = {pool.submit(request_json, url, params): name for name, (url, params) in calls.items()}
+        for future in as_completed(futures):
+            name = futures[future]
+            lines.append("")
+            lines.append(f"## {name}")
+            try:
+                items = future.result()
+                for item in items[: args.limit]:
+                    repo_id = item.get("id") or item.get("modelId") or item.get("slug") or item.get("title")
+                    likes = item.get("likes")
+                    downloads = item.get("downloads")
+                    meta = []
+                    if downloads is not None:
+                        meta.append(f"downloads={downloads}")
+                    if likes is not None:
+                        meta.append(f"likes={likes}")
+                    lines.append(f"- {repo_id}" + (f" ({', '.join(meta)})" if meta else ""))
+            except Exception as exc:
+                lines.append(f"Error: {exc}")
+    return "\n".join(lines)
+
+
+def op_snippet_search(args: argparse.Namespace) -> str:
+    if not args.query:
+        raise SystemExit("snippet_search requires --query")
+    params: dict[str, Any] = {"query": args.query, "limit": args.limit, "fields": "title,externalIds,year,citationCount"}
+    if args.date_from or args.date_to:
+        params["publicationDateOrYear"] = f"{args.date_from or ''}:{args.date_to or ''}"
+    if args.categories:
+        params["fieldsOfStudy"] = args.categories
+    if args.min_citations:
+        params["minCitationCount"] = str(args.min_citations)
+    data = request_json(f"{S2_API}/graph/v1/snippet/search", params)
+    lines = [f"# Snippet Search: {args.query}"]
+    for idx, item in enumerate(data.get("data", [])[: args.limit], 1):
+        paper = item.get("paper") or {}
+        snippet = item.get("snippet") or {}
+        lines.append(f"### {idx}. {paper.get('title', '(untitled)')}")
+        arxiv_id = paper_arxiv_id(paper)
+        if arxiv_id:
+            lines.append(f"arxiv:{arxiv_id}")
+        if snippet.get("section"):
+            lines.append(f"Section: {snippet['section']}")
+        if snippet.get("text"):
+            lines.append(f"> {truncate(snippet['text'], 400)}")
+        lines.append("")
+    return "\n".join(lines)
+
+
+def op_recommend(args: argparse.Namespace) -> str:
+    if not args.arxiv_id and not args.positive_ids:
+        raise SystemExit("recommend requires --arxiv-id or --positive-ids")
+    fields = "title,externalIds,year,citationCount,tldr,venue"
+    if args.positive_ids and not args.arxiv_id:
+        positive = [arxiv_s2_id(item.strip()) for item in args.positive_ids.split(",") if item.strip()]
+        negative = [arxiv_s2_id(item.strip()) for item in args.negative_ids.split(",") if item.strip()]
+        data = request_json(
+            f"{S2_API}/recommendations/v1/papers/",
+            {"fields": fields, "limit": args.limit},
+            method="POST",
+            body={"positivePaperIds": positive, "negativePaperIds": negative},
+        )
+    else:
+        data = request_json(
+            f"{S2_API}/recommendations/v1/papers/forpaper/{arxiv_s2_id(args.arxiv_id)}",
+            {"fields": fields, "limit": args.limit, "from": "recent"},
+        )
+    papers = data.get("recommendedPapers", [])
+    lines = ["# Recommended Papers"]
+    for idx, paper in enumerate(papers[: args.limit], 1):
+        lines.append(format_s2_paper(paper, idx))
+        lines.append("")
+    return "\n".join(lines)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("operation", choices=[
+        "trending",
+        "search",
+        "paper_details",
+        "read_paper",
+        "citation_graph",
+        "snippet_search",
+        "recommend",
+        "find_datasets",
+        "find_models",
+        "find_collections",
+        "find_all_resources",
+    ])
+    parser.add_argument("--query")
+    parser.add_argument("--arxiv-id")
+    parser.add_argument("--section")
+    parser.add_argument("--direction", choices=["citations", "references", "both"], default="both")
+    parser.add_argument("--date")
+    parser.add_argument("--date-from", default="")
+    parser.add_argument("--date-to", default="")
+    parser.add_argument("--categories")
+    parser.add_argument("--min-citations", type=int)
+    parser.add_argument("--sort-by", choices=["relevance", "citationCount", "publicationDate"], default="relevance")
+    parser.add_argument("--positive-ids", default="")
+    parser.add_argument("--negative-ids", default="")
+    parser.add_argument("--sort", choices=["downloads", "likes", "trending"], default="downloads")
+    parser.add_argument("--limit", type=int, default=10)
+    return parser
+
+
+def main() -> int:
+    args = build_parser().parse_args()
+    args.limit = min(max(args.limit, 1), 50)
+    handlers = {
+        "trending": op_trending,
+        "search": op_search,
+        "paper_details": op_paper_details,
+        "read_paper": op_read_paper,
+        "citation_graph": op_citation_graph,
+        "snippet_search": op_snippet_search,
+        "recommend": op_recommend,
+        "find_datasets": op_resources,
+        "find_models": op_resources,
+        "find_collections": op_resources,
+        "find_all_resources": op_resources,
+    }
+    try:
+        print(handlers[args.operation](args))
+    except Exception as exc:
+        print(f"Error running papers {args.operation}: {exc}", file=sys.stderr)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())