anything-llm/collector/scripts/watch/convert/as_pdf.py
Timothy Carambat 3e78476739
Franzbischoff document improvements (#241)
* cosmetic changes to be compatible to hadolint

* common configuration for most editors until better plugins comes up

* Changes on PDF metadata, using PyMuPDF (faster and more compatible)

* small changes on other file ingestions in order to try to keep the fields equal

* Lint, review, and review

* fixed unknown chars

* Use PyMuPDF for pdf loading for 200% speed increase
linting

---------

Co-authored-by: Francisco Bischoff <franzbischoff@gmail.com>
Co-authored-by: Francisco Bischoff <984592+franzbischoff@users.noreply.github.com>
2023-09-18 16:21:37 -07:00

54 lines
2.0 KiB
Python

import os, fitz
from langchain.document_loaders import PyMuPDFLoader # better UTF support and metadata
from slugify import slugify
from ..utils import guid, file_creation_time, write_to_server_documents, move_source
from ...utils import tokenize
from unidecode import unidecode
# Process all PDF-related documents.
def as_pdf(**kwargs):
parent_dir = kwargs.get('directory', 'hotdir')
filename = kwargs.get('filename')
ext = kwargs.get('ext', '.txt')
remove = kwargs.get('remove_on_complete', False)
fullpath = f"{parent_dir}/{filename}{ext}"
print(f"-- Working {fullpath} --")
loader = PyMuPDFLoader(fullpath)
pages = loader.load()
if len(pages) == 0:
print(f"{fullpath} parsing resulted in no pages - nothing to do.")
return False
# Set doc to the first page so we can still get the metadata from PyMuPDF but without all the unicode issues.
doc = pages[0]
del loader
del pages
page_content = ''
for page in fitz.open(fullpath):
print(f"-- Parsing content from pg {page.number} --")
page_content += unidecode(page.get_text('text'))
title = doc.metadata.get('title')
author = doc.metadata.get('author')
subject = doc.metadata.get('subject')
data = {
'id': guid(),
'url': "file://"+os.path.abspath(f"{parent_dir}/processed/{filename}{ext}"),
'title': title if title else f"{filename}{ext}",
'docAuthor': author if author else 'No author found',
'description': subject if subject else 'No description found.',
'docSource': 'pdf file uploaded by the user.',
'chunkSource': f"{filename}{ext}",
'published': file_creation_time(fullpath),
'wordCount': len(page_content), # Technically a letter count :p
'pageContent': page_content,
'token_count_estimate': len(tokenize(page_content))
}
write_to_server_documents(data, f"{slugify(filename)}-{data.get('id')}")
move_source(parent_dir, f"{filename}{ext}", remove=remove)
print(f"[SUCCESS]: {filename}{ext} converted & ready for embedding.\n")