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)
|