forked from Archives/langchain
Bibtex integration for document loader and retriever (#5137)
# Bibtex integration Wrap bibtexparser to retrieve a list of docs from a bibtex file. * Get the metadata from the bibtex entries * `page_content` get from the local pdf referenced in the `file` field of the bibtex entry using `pymupdf` * If no valid pdf file, `page_content` set to the `abstract` field of the bibtex entry * Support Zotero flavour using regex to get the file path * Added usage example in `docs/modules/indexes/document_loaders/examples/bibtex.ipynb` --------- Co-authored-by: Sébastien M. Popoff <sebastien.popoff@espci.fr> Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>searx_updates
Eugene Yurtsev
12 months ago
committed by
GitHub
parent
40b086d6e8
commit
5cfa72a130
@ -0,0 +1,108 @@
|
||||
import logging
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterator, List, Mapping, Optional
|
||||
|
||||
from langchain.docstore.document import Document
|
||||
from langchain.document_loaders.base import BaseLoader
|
||||
from langchain.utilities.bibtex import BibtexparserWrapper
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BibtexLoader(BaseLoader):
|
||||
"""Loads a bibtex file into a list of Documents.
|
||||
|
||||
Each document represents one entry from the bibtex file.
|
||||
|
||||
If a PDF file is present in the `file` bibtex field, the original PDF
|
||||
is loaded into the document text. If no such file entry is present,
|
||||
the `abstract` field is used instead.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
file_path: str,
|
||||
*,
|
||||
parser: Optional[BibtexparserWrapper] = None,
|
||||
max_docs: Optional[int] = None,
|
||||
max_content_chars: Optional[int] = 4_000,
|
||||
load_extra_metadata: bool = False,
|
||||
file_pattern: str = r"[^:]+\.pdf",
|
||||
):
|
||||
"""Initialize the BibtexLoader.
|
||||
|
||||
Args:
|
||||
file_path: Path to the bibtex file.
|
||||
max_docs: Max number of associated documents to load. Use -1 means
|
||||
no limit.
|
||||
"""
|
||||
self.file_path = file_path
|
||||
self.parser = parser or BibtexparserWrapper()
|
||||
self.max_docs = max_docs
|
||||
self.max_content_chars = max_content_chars
|
||||
self.load_extra_metadata = load_extra_metadata
|
||||
self.file_regex = re.compile(file_pattern)
|
||||
|
||||
def _load_entry(self, entry: Mapping[str, Any]) -> Optional[Document]:
|
||||
import fitz
|
||||
|
||||
parent_dir = Path(self.file_path).parent
|
||||
# regex is useful for Zotero flavor bibtex files
|
||||
file_names = self.file_regex.findall(entry.get("file", ""))
|
||||
if not file_names:
|
||||
return None
|
||||
texts: List[str] = []
|
||||
for file_name in file_names:
|
||||
try:
|
||||
with fitz.open(parent_dir / file_name) as f:
|
||||
texts.extend(page.get_text() for page in f)
|
||||
except FileNotFoundError as e:
|
||||
logger.debug(e)
|
||||
content = "\n".join(texts) or entry.get("abstract", "")
|
||||
if self.max_content_chars:
|
||||
content = content[: self.max_content_chars]
|
||||
metadata = self.parser.get_metadata(entry, load_extra=self.load_extra_metadata)
|
||||
return Document(
|
||||
page_content=content,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
def lazy_load(self) -> Iterator[Document]:
|
||||
"""Load bibtex file using bibtexparser and get the article texts plus the
|
||||
|
||||
article metadata.
|
||||
|
||||
See https://bibtexparser.readthedocs.io/en/master/
|
||||
|
||||
Returns:
|
||||
a list of documents with the document.page_content in text format
|
||||
"""
|
||||
try:
|
||||
import fitz # noqa: F401
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"PyMuPDF package not found, please install it with "
|
||||
"`pip install pymupdf`"
|
||||
)
|
||||
|
||||
entries = self.parser.load_bibtex_entries(self.file_path)
|
||||
if self.max_docs:
|
||||
entries = entries[: self.max_docs]
|
||||
for entry in entries:
|
||||
doc = self._load_entry(entry)
|
||||
if doc:
|
||||
yield doc
|
||||
|
||||
def load(self) -> List[Document]:
|
||||
"""Load bibtex file documents from the given bibtex file path.
|
||||
|
||||
See https://bibtexparser.readthedocs.io/en/master/
|
||||
|
||||
Args:
|
||||
file_path: the path to the bibtex file
|
||||
|
||||
Returns:
|
||||
a list of documents with the document.page_content in text format
|
||||
"""
|
||||
return list(self.lazy_load())
|
||||
@ -0,0 +1,87 @@
|
||||
"""Util that calls bibtexparser."""
|
||||
import logging
|
||||
from typing import Any, Dict, List, Mapping
|
||||
|
||||
from pydantic import BaseModel, Extra, root_validator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
OPTIONAL_FIELDS = [
|
||||
"annotate",
|
||||
"booktitle",
|
||||
"editor",
|
||||
"howpublished",
|
||||
"journal",
|
||||
"keywords",
|
||||
"note",
|
||||
"organization",
|
||||
"publisher",
|
||||
"school",
|
||||
"series",
|
||||
"type",
|
||||
"doi",
|
||||
"issn",
|
||||
"isbn",
|
||||
]
|
||||
|
||||
|
||||
class BibtexparserWrapper(BaseModel):
|
||||
"""Wrapper around bibtexparser.
|
||||
|
||||
To use, you should have the ``bibtexparser`` python package installed.
|
||||
https://bibtexparser.readthedocs.io/en/master/
|
||||
|
||||
This wrapper will use bibtexparser to load a collection of references from
|
||||
a bibtex file and fetch document summaries.
|
||||
"""
|
||||
|
||||
class Config:
|
||||
"""Configuration for this pydantic object."""
|
||||
|
||||
extra = Extra.forbid
|
||||
|
||||
@root_validator()
|
||||
def validate_environment(cls, values: Dict) -> Dict:
|
||||
"""Validate that the python package exists in environment."""
|
||||
try:
|
||||
import bibtexparser # noqa
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Could not import bibtexparser python package. "
|
||||
"Please install it with `pip install bibtexparser`."
|
||||
)
|
||||
|
||||
return values
|
||||
|
||||
def load_bibtex_entries(self, path: str) -> List[Dict[str, Any]]:
|
||||
"""Load bibtex entries from the bibtex file at the given path."""
|
||||
import bibtexparser
|
||||
|
||||
with open(path) as file:
|
||||
entries = bibtexparser.load(file).entries
|
||||
return entries
|
||||
|
||||
def get_metadata(
|
||||
self, entry: Mapping[str, Any], load_extra: bool = False
|
||||
) -> Dict[str, Any]:
|
||||
"""Get metadata for the given entry."""
|
||||
publication = entry.get("journal") or entry.get("booktitle")
|
||||
if "url" in entry:
|
||||
url = entry["url"]
|
||||
elif "doi" in entry:
|
||||
url = f'https://doi.org/{entry["doi"]}'
|
||||
else:
|
||||
url = None
|
||||
meta = {
|
||||
"id": entry.get("ID"),
|
||||
"published_year": entry.get("year"),
|
||||
"title": entry.get("title"),
|
||||
"publication": publication,
|
||||
"authors": entry.get("author"),
|
||||
"abstract": entry.get("abstract"),
|
||||
"url": url,
|
||||
}
|
||||
if load_extra:
|
||||
for field in OPTIONAL_FIELDS:
|
||||
meta[field] = entry.get(field)
|
||||
return {k: v for k, v in meta.items() if v is not None}
|
||||
@ -0,0 +1,14 @@
|
||||
@inproceedings{shen2021layoutparser,
|
||||
title = {LayoutParser: A unified toolkit for deep learning based document image analysis},
|
||||
author = {Shen, Zejiang and Zhang, Ruochen and Dell, Melissa and Lee, Benjamin Charles Germain and Carlson, Jacob and Li, Weining},
|
||||
booktitle = {Document Analysis and Recognition--ICDAR 2021: 16th International Conference, Lausanne, Switzerland, September 5--10, 2021, Proceedings, Part I 16},
|
||||
pages = {131--146},
|
||||
year = {2021},
|
||||
organization = {Springer},
|
||||
editor = {Llad{\'o}s, Josep
|
||||
and Lopresti, Daniel
|
||||
and Uchida, Seiichi},
|
||||
file = {layout-parser-paper.pdf},
|
||||
abstract = {{Recent advances in document image analysis (DIA) have been primarily driven by the application of neural networks. Ideally, research outcomes could be easily deployed in production and extended for further investigation. However, various factors like loosely organized codebases and sophisticated model configurations complicate the easy reuse of important innovations by a wide audience. Though there have been on-going efforts to improve reusability and simplify deep learning (DL) model development in disciplines like natural language processing and computer vision, none of them are optimized for challenges in the domain of DIA. This represents a major gap in the existing toolkit, as DIA is central to academic research across a wide range of disciplines in the social sciences and humanities. This paper introduces LayoutParser, an open-source library for streamlining the usage of DL in DIA research and applications. The core LayoutParser library comes with a set of simple and intuitive interfaces for applying and customizing DL models for layout detection, character recognition, and many other document processing tasks. To promote extensibility, LayoutParser also incorporates a community platform for sharing both pre-trained models and full document digitization pipelines. We demonstrate that LayoutParser is helpful for both lightweight and large-scale digitization pipelines in real-word use cases. The library is publicly available at https://layout-parser.github.io.",
|
||||
isbn="978-3-030-86549-8}},
|
||||
}
|
||||
Binary file not shown.
@ -0,0 +1,61 @@
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from langchain.document_loaders.bibtex import BibtexLoader
|
||||
|
||||
BIBTEX_EXAMPLE_FILE = Path(__file__).parent / "sample_documents" / "bibtex.bib"
|
||||
|
||||
|
||||
@pytest.mark.requires("fitz", "bibtexparser")
|
||||
def test_load_success() -> None:
|
||||
"""Test that returns one document"""
|
||||
loader = BibtexLoader(file_path=str(BIBTEX_EXAMPLE_FILE))
|
||||
docs = loader.load()
|
||||
assert len(docs) == 1
|
||||
doc = docs[0]
|
||||
assert doc.page_content
|
||||
assert set(doc.metadata) == {
|
||||
"id",
|
||||
"published_year",
|
||||
"title",
|
||||
"publication",
|
||||
"authors",
|
||||
"abstract",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.requires("fitz", "bibtexparser")
|
||||
def test_load_max_content_chars() -> None:
|
||||
"""Test that cuts off document contents at max_content_chars."""
|
||||
loader = BibtexLoader(file_path=str(BIBTEX_EXAMPLE_FILE), max_content_chars=10)
|
||||
doc = loader.load()[0]
|
||||
assert len(doc.page_content) == 10
|
||||
|
||||
|
||||
@pytest.mark.requires("fitz", "bibtexparser")
|
||||
def test_load_load_extra_metadata() -> None:
|
||||
"""Test that returns extra metadata fields."""
|
||||
loader = BibtexLoader(file_path=str(BIBTEX_EXAMPLE_FILE), load_extra_metadata=True)
|
||||
doc = loader.load()[0]
|
||||
assert set(doc.metadata) == {
|
||||
"id",
|
||||
"published_year",
|
||||
"title",
|
||||
"publication",
|
||||
"authors",
|
||||
"abstract",
|
||||
"booktitle",
|
||||
"editor",
|
||||
"organization",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.requires("fitz", "bibtexparser")
|
||||
def test_load_file_pattern() -> None:
|
||||
"""Test that returns no documents when json file pattern specified."""
|
||||
loader = BibtexLoader(
|
||||
file_path=str(BIBTEX_EXAMPLE_FILE), file_pattern=r"[^:]+\.json"
|
||||
)
|
||||
docs = loader.load()
|
||||
assert len(docs) == 0
|
||||
Loading…
Reference in New Issue