filetxt/file2text/loader.py
2025-05-12 18:13:55 +08:00

102 lines
3.3 KiB
Python

import os
import codecs
from datetime import datetime
from langchain_community.document_loaders.csv_loader import CSVLoader
from langchain_community.document_loaders.epub import UnstructuredEPubLoader
from langchain_community.document_loaders.text import TextLoader
# from langchain_community.document_loaders.pdf import UnstructuredPDFLoader
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.document_loaders.chm import UnstructuredCHMLoader
from langchain_community.document_loaders import UnstructuredWordDocumentLoader
from langchain_community.document_loaders import UnstructuredExcelLoader
from langchain_community.document_loaders import UnstructuredPowerPointLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
import mobi
import html2text
class MyMobiLoader:
def __init__(self, file_path):
self.file_path = file_path
def load(self):
tempdir, filepath = mobi.extract(self.file_path)
with codecs.open(filepath, "r", "utf-8") as f:
content=f.read()
return html2text.html2text(content)
class MyChmLoader(UnstructuredCHMLoader):
def load(self):
docs = super().load()
return ' '.join([d.page_content for d in docs])
class MyPdfLoader(PyPDFLoader):
def __init__(self, file_path, **kw):
super().__init__(file_path=file_path, extract_images=True, **kw)
def load(self):
docs = super().load()
text = ''
for d in docs:
text += ' '.join(d.page_content.split('\t'))
return text
class MyWordLoader(UnstructuredWordDocumentLoader):
def load(self):
docs = super().load()
return ' '.join([d.page_content for d in docs])
class MyPptLoader(UnstructuredPowerPointLoader):
def load(self):
docs = super().load()
return ' '.join([d.page_content for d in docs])
class MyCsvLoader(CSVLoader):
def load(self):
docs = super().load()
return ' '.join([d.page_content for d in docs])
class MyExcelLoader(UnstructuredExcelLoader):
def load(self):
docs = super().load()
return ' '.join([d.page_content for d in docs])
class MyEpubLoader(UnstructuredEPubLoader):
def load(self):
docs = super().load()
print(len(docs))
return ' '.join([d.page_content for d in docs])
def fileloader(file_path):
# Load the PDF file and split the data into chunks
data = None
if file_path.lower().endswith('.pdf'):
# loader = PyPDFLoader(file_path=file_path)
loader = MyPdfLoader(file_path=file_path)
elif file_path.lower().endswith('.docx') or file_path.lower().endswith('.doc'):
loader = MyWordLoader(file_path=file_path)
elif file_path.lower().endswith('.pptx') or file_path.lower().endswith('.pptx'):
loader = MyPptLoader(file_path=file_path)
elif file_path.lower().endswith('.xlsx') or file_path.lower().endswith('.xls'):
loader = MyExcelLoader(file_path=file_path)
elif file_path.lower().endswith('.csv'):
loader = MyCsvLoader(file_path=file_path)
elif file_path.lower().endswith('.epub'):
loader = MyEpubLoader(file_path=file_path)
elif file_path.lower().endswith('.chm'):
loader = MyChmLoader(file_path=file_path)
elif file_path.lower().endswith('.mobi'):
loader = MyMobiLoader(file_path=file_path)
else:
loader = TextLoader(file_path=file_path)
data = loader.load()
return data
if __name__ == '__main__':
import sys
if len(sys.argv) < 2:
print(f'{sys.argv[0]} file\nload a file and get its text')
sys.exit(1)
text = fileloader(sys.argv[1])
print(f'{text}')