filetxt/file2text/loader.py
2025-05-13 14:20:13 +08:00

155 lines
5.1 KiB
Python

import os
import codecs
from datetime import datetime
from langchain_community.document_loaders.csv_loader import CSVLoader
from langchain_community.document_loaders.text import TextLoader
# from langchain_community.document_loaders.epub import UnstructuredEPubLoader
# from langchain_community.document_loaders.pdf import UnstructuredPDFLoader
# from langchain_community.document_loaders import PyPDFLoader
# from langchain_community.document_loaders import UnstructuredWordDocumentLoader
from langchain_community.document_loaders.chm import UnstructuredCHMLoader
from langchain_community.document_loaders import UnstructuredExcelLoader
from langchain_community.document_loaders import UnstructuredPowerPointLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from docx import Document
# import chm
from pptx import Presentation
from openpyxl import load_workbook
import mobi
import PyPDF2
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 __init__(self, file_path):
self.filepath = file_path
self.chm_file = None
def load(self):
"""Reads the CHM file and returns concatenated text from all text/html entries."""
self.chm_file = chm.CHMFile(self.filepath)
content = []
def callback(chm_item):
if chm_item[0].endswith(('.html', '.htm', '.txt')):
try:
data = self.chm_file.read_file(chm_item[0])
if data:
content.append(data.decode('utf-8', errors='ignore'))
except Exception:
pass
return True
self.chm_file.walk(callback)
return '\n'.join(content)
class MyPdfLoader:
def __init__(self, file_path, **kw):
self.filepath = file_path
self.reader = None
def load(self):
"""Reads the PDF file and returns all text as a single string."""
text = ''
with open(self.filepath, 'rb') as file:
self.reader = PyPDF2.PdfReader(file)
for page in self.reader.pages:
text += page.extract_text() or ''
return ' '.join(text.split('\t'))
class MyWordLoader:
def __init__(self, file_path):
self.filepath = file_path
self.document = None
def load(self):
"""Reads the .docx file and returns the full text as a single string."""
self.document = Document(self.filepath)
text = '\n'.join([para.text for para in self.document.paragraphs])
return text
class MyPptLoader:
def __init__(self, file_path):
self.filepath = file_path
def load(self):
prs = Presentation(self.filepath)
text = []
for slide in prs.slides:
for shape in slide.shapes:
if hasattr(shape, "text"):
text.append(shape.text)
return "\n".join(text)
class MyCsvLoader(CSVLoader):
def load(self):
docs = super().load()
return ' '.join([d.page_content for d in docs])
class MyExcelLoader:
def __init__(self, file_path):
self.filepath = file_path
self.workbook = None
def load(self):
"""Reads all sheets in the Excel file and returns the content as a string."""
self.workbook = load_workbook(filename=self.filepath, data_only=True)
content = []
for sheet in self.workbook.worksheets:
content.append(f"--- Sheet: {sheet.title} ---")
for row in sheet.iter_rows(values_only=True):
row_text = '\t'.join(str(cell) if cell is not None else '' for cell in row)
content.append(row_text)
return '\n'.join(content)
class MyEpubLoader:
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}')