lidp
commited on
Commit
·
edccacf
1
Parent(s):
f2ca0b5
Fix docx images (#2756)
Browse files### What problem does this PR solve?
#2755
### Type of change
- [x] Bug Fix (non-breaking change which fixes an issue)
- rag/app/naive.py +18 -9
rag/app/naive.py
CHANGED
|
@@ -16,14 +16,15 @@ from docx import Document
|
|
| 16 |
from timeit import default_timer as timer
|
| 17 |
import re
|
| 18 |
from deepdoc.parser.pdf_parser import PlainParser
|
| 19 |
-
from rag.nlp import rag_tokenizer, naive_merge, tokenize_table, tokenize_chunks, find_codec, concat_img,
|
|
|
|
| 20 |
from deepdoc.parser import PdfParser, ExcelParser, DocxParser, HtmlParser, JsonParser, MarkdownParser, TxtParser
|
| 21 |
from rag.settings import cron_logger
|
| 22 |
from rag.utils import num_tokens_from_string
|
| 23 |
from PIL import Image
|
| 24 |
from functools import reduce
|
| 25 |
from markdown import markdown
|
| 26 |
-
from docx.image.exceptions import UnrecognizedImageError
|
| 27 |
|
| 28 |
|
| 29 |
class Docx(DocxParser):
|
|
@@ -42,6 +43,12 @@ class Docx(DocxParser):
|
|
| 42 |
except UnrecognizedImageError:
|
| 43 |
print("Unrecognized image format. Skipping image.")
|
| 44 |
return None
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 45 |
try:
|
| 46 |
image = Image.open(BytesIO(image_blob)).convert('RGB')
|
| 47 |
return image
|
|
@@ -101,7 +108,7 @@ class Docx(DocxParser):
|
|
| 101 |
while i < len(r.cells):
|
| 102 |
span = 1
|
| 103 |
c = r.cells[i]
|
| 104 |
-
for j in range(i+1, len(r.cells)):
|
| 105 |
if c.text == r.cells[j].text:
|
| 106 |
span += 1
|
| 107 |
i = j
|
|
@@ -136,9 +143,9 @@ class Pdf(PdfParser):
|
|
| 136 |
self._text_merge()
|
| 137 |
callback(0.67, "Text merging finished")
|
| 138 |
tbls = self._extract_table_figure(True, zoomin, True, True)
|
| 139 |
-
#self._naive_vertical_merge()
|
| 140 |
self._concat_downward()
|
| 141 |
-
#self._filter_forpages()
|
| 142 |
|
| 143 |
cron_logger.info("layouts: {}".format(timer() - start))
|
| 144 |
return [(b["text"], self._line_tag(b, zoomin))
|
|
@@ -158,8 +165,8 @@ class Markdown(MarkdownParser):
|
|
| 158 |
tbls = []
|
| 159 |
for sec in remainder.split("\n"):
|
| 160 |
if num_tokens_from_string(sec) > 10 * self.chunk_token_num:
|
| 161 |
-
sections.append((sec[:int(len(sec)/2)], ""))
|
| 162 |
-
sections.append((sec[int(len(sec)/2):], ""))
|
| 163 |
else:
|
| 164 |
sections.append((sec, ""))
|
| 165 |
print(tables)
|
|
@@ -191,7 +198,7 @@ def chunk(filename, binary=None, from_page=0, to_page=100000,
|
|
| 191 |
if re.search(r"\.docx$", filename, re.IGNORECASE):
|
| 192 |
callback(0.1, "Start to parse.")
|
| 193 |
sections, tbls = Docx()(filename, binary)
|
| 194 |
-
res = tokenize_table(tbls, doc, eng)
|
| 195 |
|
| 196 |
callback(0.8, "Finish parsing.")
|
| 197 |
st = timer()
|
|
@@ -229,7 +236,7 @@ def chunk(filename, binary=None, from_page=0, to_page=100000,
|
|
| 229 |
parser_config.get("chunk_token_num", 128),
|
| 230 |
parser_config.get("delimiter", "\n!?;。;!?"))
|
| 231 |
callback(0.8, "Finish parsing.")
|
| 232 |
-
|
| 233 |
elif re.search(r"\.(md|markdown)$", filename, re.IGNORECASE):
|
| 234 |
callback(0.1, "Start to parse.")
|
| 235 |
sections, tbls = Markdown(int(parser_config.get("chunk_token_num", 128)))(filename, binary)
|
|
@@ -276,7 +283,9 @@ def chunk(filename, binary=None, from_page=0, to_page=100000,
|
|
| 276 |
if __name__ == "__main__":
|
| 277 |
import sys
|
| 278 |
|
|
|
|
| 279 |
def dummy(prog=None, msg=""):
|
| 280 |
pass
|
| 281 |
|
|
|
|
| 282 |
chunk(sys.argv[1], from_page=0, to_page=10, callback=dummy)
|
|
|
|
| 16 |
from timeit import default_timer as timer
|
| 17 |
import re
|
| 18 |
from deepdoc.parser.pdf_parser import PlainParser
|
| 19 |
+
from rag.nlp import rag_tokenizer, naive_merge, tokenize_table, tokenize_chunks, find_codec, concat_img, \
|
| 20 |
+
naive_merge_docx, tokenize_chunks_docx
|
| 21 |
from deepdoc.parser import PdfParser, ExcelParser, DocxParser, HtmlParser, JsonParser, MarkdownParser, TxtParser
|
| 22 |
from rag.settings import cron_logger
|
| 23 |
from rag.utils import num_tokens_from_string
|
| 24 |
from PIL import Image
|
| 25 |
from functools import reduce
|
| 26 |
from markdown import markdown
|
| 27 |
+
from docx.image.exceptions import UnrecognizedImageError, UnexpectedEndOfFileError, InvalidImageStreamError
|
| 28 |
|
| 29 |
|
| 30 |
class Docx(DocxParser):
|
|
|
|
| 43 |
except UnrecognizedImageError:
|
| 44 |
print("Unrecognized image format. Skipping image.")
|
| 45 |
return None
|
| 46 |
+
except UnexpectedEndOfFileError:
|
| 47 |
+
print("EOF was unexpectedly encountered while reading an image stream. Skipping image.")
|
| 48 |
+
return None
|
| 49 |
+
except InvalidImageStreamError:
|
| 50 |
+
print("The recognized image stream appears to be corrupted. Skipping image.")
|
| 51 |
+
return None
|
| 52 |
try:
|
| 53 |
image = Image.open(BytesIO(image_blob)).convert('RGB')
|
| 54 |
return image
|
|
|
|
| 108 |
while i < len(r.cells):
|
| 109 |
span = 1
|
| 110 |
c = r.cells[i]
|
| 111 |
+
for j in range(i + 1, len(r.cells)):
|
| 112 |
if c.text == r.cells[j].text:
|
| 113 |
span += 1
|
| 114 |
i = j
|
|
|
|
| 143 |
self._text_merge()
|
| 144 |
callback(0.67, "Text merging finished")
|
| 145 |
tbls = self._extract_table_figure(True, zoomin, True, True)
|
| 146 |
+
# self._naive_vertical_merge()
|
| 147 |
self._concat_downward()
|
| 148 |
+
# self._filter_forpages()
|
| 149 |
|
| 150 |
cron_logger.info("layouts: {}".format(timer() - start))
|
| 151 |
return [(b["text"], self._line_tag(b, zoomin))
|
|
|
|
| 165 |
tbls = []
|
| 166 |
for sec in remainder.split("\n"):
|
| 167 |
if num_tokens_from_string(sec) > 10 * self.chunk_token_num:
|
| 168 |
+
sections.append((sec[:int(len(sec) / 2)], ""))
|
| 169 |
+
sections.append((sec[int(len(sec) / 2):], ""))
|
| 170 |
else:
|
| 171 |
sections.append((sec, ""))
|
| 172 |
print(tables)
|
|
|
|
| 198 |
if re.search(r"\.docx$", filename, re.IGNORECASE):
|
| 199 |
callback(0.1, "Start to parse.")
|
| 200 |
sections, tbls = Docx()(filename, binary)
|
| 201 |
+
res = tokenize_table(tbls, doc, eng) # just for table
|
| 202 |
|
| 203 |
callback(0.8, "Finish parsing.")
|
| 204 |
st = timer()
|
|
|
|
| 236 |
parser_config.get("chunk_token_num", 128),
|
| 237 |
parser_config.get("delimiter", "\n!?;。;!?"))
|
| 238 |
callback(0.8, "Finish parsing.")
|
| 239 |
+
|
| 240 |
elif re.search(r"\.(md|markdown)$", filename, re.IGNORECASE):
|
| 241 |
callback(0.1, "Start to parse.")
|
| 242 |
sections, tbls = Markdown(int(parser_config.get("chunk_token_num", 128)))(filename, binary)
|
|
|
|
| 283 |
if __name__ == "__main__":
|
| 284 |
import sys
|
| 285 |
|
| 286 |
+
|
| 287 |
def dummy(prog=None, msg=""):
|
| 288 |
pass
|
| 289 |
|
| 290 |
+
|
| 291 |
chunk(sys.argv[1], from_page=0, to_page=10, callback=dummy)
|