Merge pull request #209 from TideDra/dev

TideDra · web-flow · commit 7b18a4b2ce9e · 2026-03-25T18:09:30.000+08:00
extract text from arxiv html
diff --git a/README.md b/README.md
@@ -154,7 +154,6 @@ reranker:
 executor:
   debug: false # Whether to use debug mode. Example: true
   send_empty: false # Whether to send an empty email even if no new papers today. Example: true
-  max_workers: 4 # Concurrent workers for processing papers. Example: 4
   max_paper_num: 100 # The maximum number of the papers presented in the email. Example: 100
   source: ??? # The sources of papers to retrieve. Example: ['arxiv','biorxiv','medrxiv']
   reranker: local # The reranker to use. Example: 'local' or 'api'
diff --git a/config/base.yaml b/config/base.yaml
@@ -45,7 +45,6 @@ reranker:
 executor:
   debug: false # Whether to use debug mode. Example: true
   send_empty: false # Whether to send an empty email even if no new papers today. Example: true
-  max_workers: 4 # Concurrent workers for processing papers. Example: 4
   max_paper_num: 100 # The maximum number of the papers presented in the email. Example: 100
   source: ??? # The sources of papers to retrieve. Example: ['arxiv','biorxiv','medrxiv']
   reranker: local # The reranker to use. Example: 'local' or 'api'
diff --git a/pyproject.toml b/pyproject.toml
@@ -25,6 +25,7 @@ dependencies = [
     "pymupdf-layout>=1.27.1",
     "dotenv>=0.9.9",
     "peft>=0.18.1",
+    "trafilatura>=2.0.0",
 ]
 
 [tool.uv.sources]
diff --git a/src/zotero_arxiv_daily/retriever/arxiv_retriever.py b/src/zotero_arxiv_daily/retriever/arxiv_retriever.py
@@ -4,22 +4,116 @@
 from ..protocol import Paper
 from ..utils import extract_markdown_from_pdf, extract_tex_code_from_tar
 from tempfile import TemporaryDirectory
-from concurrent.futures import ThreadPoolExecutor, TimeoutError
 import feedparser
-from urllib.request import urlretrieve
 from tqdm import tqdm
+import multiprocessing
 import os
+from queue import Empty
+from typing import Any, Callable, TypeVar
 from loguru import logger
+import requests
 
+T = TypeVar("T")
+
+DOWNLOAD_TIMEOUT = (10, 60)
 PDF_EXTRACT_TIMEOUT = 180
+TAR_EXTRACT_TIMEOUT = 180
+
+
+def _download_file(url: str, path: str) -> None:
+    with requests.get(url, stream=True, timeout=DOWNLOAD_TIMEOUT) as response:
+        response.raise_for_status()
+        with open(path, "wb") as file:
+            for chunk in response.iter_content(chunk_size=1024 * 1024):
+                if chunk:
+                    file.write(chunk)
+
+
+def _run_in_subprocess(
+    result_queue: Any,
+    func: Callable[..., T | None],
+    args: tuple[Any, ...],
+) -> None:
+    try:
+        result_queue.put(("ok", func(*args)))
+    except Exception as exc:
+        result_queue.put(("error", f"{type(exc).__name__}: {exc}"))
+
+
+def _run_with_hard_timeout(
+    func: Callable[..., T | None],
+    args: tuple[Any, ...],
+    *,
+    timeout: float,
+    operation: str,
+    paper_title: str,
+) -> T | None:
+    start_methods = multiprocessing.get_all_start_methods()
+    context = multiprocessing.get_context("fork" if "fork" in start_methods else start_methods[0])
+    result_queue = context.Queue()
+    process = context.Process(target=_run_in_subprocess, args=(result_queue, func, args))
+    process.start()
+
+    try:
+        status, payload = result_queue.get(timeout=timeout)
+    except Empty:
+        if process.is_alive():
+            process.kill()
+        process.join(5)
+        result_queue.close()
+        result_queue.join_thread()
+        logger.warning(f"{operation} timed out for {paper_title} after {timeout} seconds")
+        return None
+
+    process.join(5)
+    result_queue.close()
+    result_queue.join_thread()
+
+    if status == "ok":
+        return payload
+
+    logger.warning(f"{operation} failed for {paper_title}: {payload}")
+    return None
+
+
+def _extract_text_from_pdf_worker(pdf_url: str) -> str:
+    with TemporaryDirectory() as temp_dir:
+        path = os.path.join(temp_dir, "paper.pdf")
+        _download_file(pdf_url, path)
+        return extract_markdown_from_pdf(path)
+
+
+def _extract_text_from_html_worker(html_url: str) -> str | None:
+    import trafilatura
+
+    downloaded = trafilatura.fetch_url(html_url)
+    if downloaded is None:
+        raise ValueError(f"Failed to download HTML from {html_url}")
+    text = trafilatura.extract(downloaded, include_comments=False, include_tables=False)
+    if not text:
+        raise ValueError(f"No text extracted from {html_url}")
+    return text
+
+
+def _extract_text_from_tar_worker(source_url: str, paper_id: str) -> str | None:
+    with TemporaryDirectory() as temp_dir:
+        path = os.path.join(temp_dir, "paper.tar.gz")
+        _download_file(source_url, path)
+        file_contents = extract_tex_code_from_tar(path, paper_id)
+        if not file_contents or "all" not in file_contents:
+            raise ValueError("Main tex file not found.")
+        return file_contents["all"]
+
+
 @register_retriever("arxiv")
 class ArxivRetriever(BaseRetriever):
     def __init__(self, config):
         super().__init__(config)
         if self.config.source.arxiv.category is None:
             raise ValueError("category must be specified for arxiv.")
+
     def _retrieve_raw_papers(self) -> list[ArxivResult]:
-        client = arxiv.Client(num_retries=10,delay_seconds=10)
+        client = arxiv.Client(num_retries=10, delay_seconds=10)
         query = '+'.join(self.config.source.arxiv.category)
         include_cross_list = self.config.source.arxiv.get("include_cross_list", False)
         # Get the latest paper from arxiv rss feed
@@ -38,26 +132,23 @@ def _retrieve_raw_papers(self) -> list[ArxivResult]:
 
         # Get full information of each paper from arxiv api
         bar = tqdm(total=len(all_paper_ids))
-        for i in range(0,len(all_paper_ids),20):
-            search = arxiv.Search(id_list=all_paper_ids[i:i+20])
+        for i in range(0, len(all_paper_ids), 20):
+            search = arxiv.Search(id_list=all_paper_ids[i:i + 20])
             batch = list(client.results(search))
             bar.update(len(batch))
             raw_papers.extend(batch)
         bar.close()
 
         return raw_papers
 
-    def convert_to_paper(self, raw_paper:ArxivResult) -> Paper:
+    def convert_to_paper(self, raw_paper: ArxivResult) -> Paper:
         title = raw_paper.title
         authors = [a.name for a in raw_paper.authors]
         abstract = raw_paper.summary
         pdf_url = raw_paper.pdf_url
-        try:
-            with ThreadPoolExecutor(max_workers=1) as pool:
-                full_text = pool.submit(extract_text_from_pdf, raw_paper).result(timeout=PDF_EXTRACT_TIMEOUT)
-        except TimeoutError:
-            logger.warning(f"PDF extraction timed out for {raw_paper.title}")
-            full_text = None
+        full_text = extract_text_from_html(raw_paper)
+        if full_text is None:
+            full_text = extract_text_from_pdf(raw_paper)
         if full_text is None:
             full_text = extract_text_from_tar(raw_paper)
         return Paper(
@@ -67,46 +158,41 @@ def convert_to_paper(self, raw_paper:ArxivResult) -> Paper:
             abstract=abstract,
             url=raw_paper.entry_id,
             pdf_url=pdf_url,
-            full_text=full_text
+            full_text=full_text,
         )
 
+
+def extract_text_from_html(paper: ArxivResult) -> str | None:
+    html_url = paper.entry_id.replace("/abs/", "/html/")
+    try:
+        return _extract_text_from_html_worker(html_url)
+    except Exception as exc:
+        logger.warning(f"HTML extraction failed for {paper.title}: {exc}")
+        return None
+
+
 def extract_text_from_pdf(paper: ArxivResult) -> str | None:
-    with TemporaryDirectory() as temp_dir:
-        path = os.path.join(temp_dir, "paper.pdf")
-        if paper.pdf_url is None:
-            logger.warning(f"No PDF URL available for {paper.title}")
-            return None
-        try:
-            urlretrieve(paper.pdf_url, path)
-        except Exception as e:
-            logger.warning(f"Failed to download pdf for {paper.title}: {type(e).__name__}: {e}")
-            return None
-        try:
-            full_text = extract_markdown_from_pdf(path)
-        except Exception as e:
-            logger.warning(f"Failed to extract full text of {paper.title} from pdf: {e}")
-            full_text = None
-        return full_text
+    if paper.pdf_url is None:
+        logger.warning(f"No PDF URL available for {paper.title}")
+        return None
+    return _run_with_hard_timeout(
+        _extract_text_from_pdf_worker,
+        (paper.pdf_url,),
+        timeout=PDF_EXTRACT_TIMEOUT,
+        operation="PDF extraction",
+        paper_title=paper.title,
+    )
+
 
 def extract_text_from_tar(paper: ArxivResult) -> str | None:
-    with TemporaryDirectory() as temp_dir:
-        path = os.path.join(temp_dir, "paper.tar.gz")
-        source_url = paper.source_url()
-        if source_url is None:
-            logger.warning(f"No source URL available for {paper.title}")
-            return None
-        try:
-            urlretrieve(source_url, path)
-        except Exception as e:
-            logger.warning(f"Failed to download source for {paper.title}: {type(e).__name__}: {e}")
-            return None
-        try:
-            file_contents = extract_tex_code_from_tar(path, paper.entry_id)
-            if "all" not in file_contents:
-                logger.warning(f"Failed to extract full text of {paper.title} from tar: Main tex file not found.")
-                return None
-            full_text = file_contents["all"]
-        except Exception as e:
-            logger.warning(f"Failed to extract full text of {paper.title} from tar: {e}")
-            full_text = None
-        return full_text
+    source_url = paper.source_url()
+    if source_url is None:
+        logger.warning(f"No source URL available for {paper.title}")
+        return None
+    return _run_with_hard_timeout(
+        _extract_text_from_tar_worker,
+        (source_url, paper.entry_id),
+        timeout=TAR_EXTRACT_TIMEOUT,
+        operation="Tar extraction",
+        paper_title=paper.title,
+    )
diff --git a/src/zotero_arxiv_daily/retriever/base.py b/src/zotero_arxiv_daily/retriever/base.py
@@ -1,34 +1,12 @@
 from abc import ABC, abstractmethod
 from omegaconf import DictConfig
 from ..protocol import Paper, RawPaperItem
-from concurrent.futures import ProcessPoolExecutor, as_completed
 from tqdm import tqdm
 from typing import Type
+from time import sleep
 from loguru import logger
 
 
-def _describe_raw_paper(raw_paper: RawPaperItem) -> str:
-    title = getattr(raw_paper, "title", None)
-    if title:
-        return str(title)
-    if isinstance(raw_paper, dict):
-        for key in ("title", "entry_id", "id", "doi"):
-            value = raw_paper.get(key)
-            if value:
-                return str(value)
-    return repr(raw_paper)
-
-
-def _convert_to_paper_safe(retriever: "BaseRetriever", raw_paper: RawPaperItem) -> Paper | None:
-    try:
-        return retriever.convert_to_paper(raw_paper)
-    except Exception as exc:
-        logger.warning(
-            f"Skipping paper {_describe_raw_paper(raw_paper)}: {type(exc).__name__}: {exc}"
-        )
-        return None
-
-
 class BaseRetriever(ABC):
     name: str
     def __init__(self, config:DictConfig):
@@ -45,21 +23,18 @@ def convert_to_paper(self, raw_paper:RawPaperItem) -> Paper | None:
 
     def retrieve_papers(self) -> list[Paper]:
         raw_papers = self._retrieve_raw_papers()
-        papers = []
         logger.info("Processing papers...")
-        with ProcessPoolExecutor(max_workers=self.config.executor.max_workers) as exec_pool:
-            futures = {exec_pool.submit(_convert_to_paper_safe, self, rp): i for i, rp in enumerate(raw_papers)}
-            papers = [None] * len(raw_papers)
-            for future in tqdm(as_completed(futures), total=len(raw_papers), desc="Converting papers"):
-                try:
-                    papers[futures[future]] = future.result()
-                except Exception as exc:
-                    raw_paper = raw_papers[futures[future]]
-                    logger.warning(
-                        f"Skipping paper {_describe_raw_paper(raw_paper)} after worker failure: "
-                        f"{type(exc).__name__}: {exc}"
-                    )
-        return [p for p in papers if p is not None]
+        papers = []
+        for raw_paper in tqdm(raw_papers, total=len(raw_papers), desc="Converting papers"):
+            try:
+                paper = self.convert_to_paper(raw_paper)
+            except Exception as exc:
+                logger.warning(f"Skipping paper {getattr(raw_paper, 'title', raw_paper)}: {exc}")
+                continue
+            if paper is not None:
+                papers.append(paper)
+            sleep(1)
+        return papers
 
 registered_retrievers = {}
 
diff --git a/tests/retriever/test_arxiv_retriever.py b/tests/retriever/test_arxiv_retriever.py
diff --git a/uv.lock b/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -25,6 +25,7 @@ dependencies = [`
`25`	`25`	`"pymupdf-layout>=1.27.1",`
`26`	`26`	`"dotenv>=0.9.9",`
`27`	`27`	`"peft>=0.18.1",`
	`28`	`+ "trafilatura>=2.0.0",`
`28`	`29`	`]`
`29`	`30`
`30`	`31`	`[tool.uv.sources]`