102 lines
3.3 KiB
Python
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}')
|
|
|