File size: 7,628 Bytes
42da440 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 | import re
import io
import chardet
import os
try:
from PIL import Image
from pdf2image import convert_from_path
import pytesseract
import fitz # pymupdf
import docx
def extract_text_from_docx(path):
"""
Extracts text from a DOCX file.
Args:
path (str): The file path to the DOCX file.
Returns:
str: The extracted text.
"""
doc = docx.Document(path)
full_text = []
for para in doc.paragraphs:
full_text.append(para.text)
return "\n".join(full_text)
def extract_text_from_pdf(path):
"""
Extracts text from a copyable PDF using PyMuPDF.
Args:
path (str): The file path to the PDF.
Returns:
str: The extracted text.
"""
text = ""
with fitz.open(path) as doc:
for page in doc:
text += page.get_text()
return text
def extract_text_from_pdf_ocr(path):
"""
Extracts text from a non-copyable PDF using OCR.
Args:
path (str): The file path to the PDF.
Returns:
str: The extracted text.
"""
text = ""
with fitz.open(path) as doc:
for page_number, page in enumerate(doc):
# logger.info("Performing OCR on page %d", page_number + 1)
pix = page.get_pixmap()
img_bytes = pix.tobytes("png")
img = Image.open(io.BytesIO(img_bytes))
page_text = pytesseract.image_to_string(img)
text += page_text + "\n"
return text
def remove_newlines_in_sentences(text):
lines = text.split("\n")
new_lines = []
for line in lines:
line = line.strip()
if line:
if line[-1] not in ".!?":
line += " "
else:
line += "\n"
new_lines.append(line)
return "".join(new_lines)
def extract_text(path):
"""
Extracts formatted text from a PDF or DOCX file.
Args:
path (str): The file path to the PDF or DOCX file.
Returns:
str: The extracted text in markdown format.
"""
# Check the file extension
_, ext = os.path.splitext(path)
ext = ext.lower()
if ext == ".docx":
# logger.info("Extracting text from DOCX file.")
text = extract_text_from_docx(path)
elif ext == ".pdf":
# logger.info("Extracting text from PDF file.")
text = extract_text_from_pdf(path)
# logger.info("Extracted text length: %d", len(text))
# If extracted text is too short, use OCR
if len(text.strip()) < 100:
# logger.info("Extracted text is too short, switching to OCR.")
text = extract_text_from_pdf_ocr(path)
else:
raise ValueError(f"Unsupported file type: {ext}")
# Remove newlines within sentences
text = remove_newlines_in_sentences(text)
return text
except ImportError as e:
print(
"NOTE PDF and DOCX extraction will not work without the required libraries. Please install the required libraries to enable this functionality. You need: `PyMuPDF Pillow pytesseract python-docx`. You also need to have Tesseract OCR installed on your system. Visit: https://tesseract-ocr.github.io/tessdoc/Installation.html (these are not included by default because tesseract can be a pain to setup)"
)
print("This is the error")
print(e)
def extract_text(file_path):
raise Exception(
"PDF and DOCX extraction is not supported without the required libraries."
)
def sentence_chunking_algorithm(file_path, max_char_length=1900):
"""
This function takes a plaintext file and chunks it into paragraphs or sentences if the paragraph exceeds max_char_length.
:param file_path: Path to the plaintext file
:param max_char_length: The maximum char5acter length for a chunk
:return: List of chunks with source text information
"""
chunks_with_source = []
current_chunk = []
char_count = 0
source_name = os.path.basename(file_path)
if file_path.endswith(".pdf") or file_path.endswith(".docx"):
content = extract_text(file_path)
else:
# with open(file_path, 'rb') as raw_file:
# raw_data = raw_file.read()
# result = chardet.detect(raw_data)
# file_encoding = result['encoding']
# Now read the file with the detected encoding
with open(file_path, "r", errors="ignore") as file:
content = file.read()
# with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
# content = f.read()
# try:
# with open(file_path, "r", encoding="utf-8") as f:
# content = f.read()
# except Exception as e:
# print(f"\nError reading file {file_path}: {e}\n")
# return []
paragraphs = content.split("\n\n")
max_char_length = int(max_char_length)
for paragraph in paragraphs:
paragraph = paragraph.strip() # Remove leading and trailing whitespace
if not paragraph: # Skip empty paragraphs
continue
paragraph_char_count = len(paragraph)
# Check if the paragraph itself exceeds the max token length
if paragraph_char_count > max_char_length:
# Fallback to character chunking for this paragraph
end_index = 0
while end_index < paragraph_char_count:
chunk_end = min(end_index + max_char_length, paragraph_char_count)
# Take until the next sentence ends (or we reach max_char_length*1.5)
while (
chunk_end < paragraph_char_count
and paragraph[chunk_end] not in [".", "!", "?", "\n"]
and chunk_end < end_index + max_char_length * 1.5
):
chunk_end += 1
# add one to chunk_end to include the punctuation IF it's not the last character
if chunk_end < paragraph_char_count:
chunk_end += 1
chunks_with_source.append(
{
"paragraph": paragraph[end_index:chunk_end],
"metadata": source_name,
}
)
end_index = chunk_end
# # handle the remainder of the paragraph
# end_index = end_index - max_char_length
# current_chunk.append(paragraph[end_index:])
# char_count = paragraph_char_count - end_index
else:
if char_count + paragraph_char_count <= max_char_length:
current_chunk.append(paragraph)
char_count += paragraph_char_count
else:
chunks_with_source.append(
{"paragraph": "".join(current_chunk), "metadata": source_name}
)
current_chunk = [paragraph]
char_count = paragraph_char_count
# Add the last chunk if it exists
if current_chunk:
chunks_with_source.append(
{"paragraph": "".join(current_chunk), "metadata": source_name}
)
# filter out chunks with fewer than 50 characters
chunks_with_source = [
chunk for chunk in chunks_with_source if len(chunk["paragraph"]) >= 50
]
return chunks_with_source, content
|