Merge branch 'dev' of https://github.com/imClumsyPanda/langchain-ChatGLM into dev

chains/local_doc_qa.py

@@ -10,6 +10,8 @@ import numpy as np
 from utils import torch_gc
 from tqdm import tqdm
 from pypinyin import lazy_pinyin
+from loader import UnstructuredPaddleImageLoader
+from loader import UnstructuredPaddlePDFLoader
 
 DEVICE_ = EMBEDDING_DEVICE
 DEVICE_ID = "0" if torch.cuda.is_available() else None
@@ -21,16 +23,31 @@ def load_file(filepath, sentence_size=SENTENCE_SIZE):
         loader = UnstructuredFileLoader(filepath, mode="elements")
         docs = loader.load()
     elif filepath.lower().endswith(".pdf"):
-        loader = UnstructuredFileLoader(filepath, strategy="fast")
+        loader = UnstructuredPaddlePDFLoader(filepath)
         textsplitter = ChineseTextSplitter(pdf=True, sentence_size=sentence_size)
         docs = loader.load_and_split(textsplitter)
+    elif filepath.lower().endswith(".jpg") or filepath.lower().endswith(".png"):
+        loader = UnstructuredPaddleImageLoader(filepath, mode="elements")
+        textsplitter = ChineseTextSplitter(pdf=False, sentence_size=sentence_size)
+        docs = loader.load_and_split(text_splitter=textsplitter)
     else:
         loader = UnstructuredFileLoader(filepath, mode="elements")
         textsplitter = ChineseTextSplitter(pdf=False, sentence_size=sentence_size)
         docs = loader.load_and_split(text_splitter=textsplitter)
+    write_check_file(filepath, docs)
     return docs
 
 
+def write_check_file(filepath, docs):
+    fout = open('load_file.txt', 'a')
+    fout.write("filepath=%s,len=%s" % (filepath, len(docs)))
+    fout.write('\n')
+    for i in docs:
+        fout.write(str(i))
+        fout.write('\n')
+    fout.close()
+
+
 def generate_prompt(related_docs: List[str], query: str,
                     prompt_template=PROMPT_TEMPLATE) -> str:
     context = "\n".join([doc.page_content for doc in related_docs])

loader/__init__.py

@@ -0,0 +1,2 @@
+from .image_loader import UnstructuredPaddleImageLoader
+from .pdf_loader import UnstructuredPaddlePDFLoader

loader/image_loader.py

@@ -0,0 +1,28 @@
+"""Loader that loads image files."""
+from typing import List
+
+from langchain.document_loaders.unstructured import UnstructuredFileLoader
+from paddleocr import PaddleOCR
+import os
+
+
+class UnstructuredPaddleImageLoader(UnstructuredFileLoader):
+    """Loader that uses unstructured to load image files, such as PNGs and JPGs."""
+
+    def _get_elements(self) -> List:
+        def image_ocr_txt(filepath, dir_path="tmp_files"):
+            if not os.path.exists(dir_path):
+                os.makedirs(dir_path)
+            filename = os.path.split(filepath)[-1]
+            ocr = PaddleOCR(lang="ch", use_gpu=False, show_log=False)
+            result = ocr.ocr(img=filepath)
+
+            ocr_result = [i[1][0] for line in result for i in line]
+            txt_file_path = os.path.join(dir_path, "%s.txt" % (filename))
+            with open(txt_file_path, 'w', encoding='utf-8') as fout:
+                fout.write("\n".join(ocr_result))
+            return txt_file_path
+
+        txt_file_path = image_ocr_txt(self.file_path)
+        from unstructured.partition.text import partition_text
+        return partition_text(filename=txt_file_path, **self.unstructured_kwargs)

loader/pdf_loader.py

@@ -0,0 +1,44 @@
+"""Loader that loads image files."""
+from typing import List
+
+from langchain.document_loaders.unstructured import UnstructuredFileLoader
+from paddleocr import PaddleOCR
+import os
+import fitz
+
+
+class UnstructuredPaddlePDFLoader(UnstructuredFileLoader):
+    """Loader that uses unstructured to load image files, such as PNGs and JPGs."""
+
+    def _get_elements(self) -> List:
+        def pdf_ocr_txt(filepath, dir_path="tmp_files"):
+            if not os.path.exists(dir_path):
+                os.makedirs(dir_path)
+            filename = os.path.split(filepath)[-1]
+            ocr = PaddleOCR(lang="ch", use_gpu=False, show_log=False)
+            doc = fitz.open(filepath)
+            txt_file_path = os.path.join(dir_path, "%s.txt" % (filename))
+            img_name = './img/.tmp.png'
+            with open(txt_file_path, 'w', encoding='utf-8') as fout:
+
+                for i in range(doc.page_count):
+                    page = doc[i]
+                    text = page.get_text("")
+                    fout.write(text)
+                    fout.write("\n")
+
+                    img_list = page.get_images()
+                    for img in img_list:
+                        pix = fitz.Pixmap(doc, img[0])
+
+                        pix.save(img_name)
+
+                        result = ocr.ocr(img_name)
+                        ocr_result = [i[1][0] for line in result for i in line]
+                        fout.write("\n".join(ocr_result))
+            os.remove(img_name)
+            return txt_file_path
+
+        txt_file_path = pdf_ocr_txt(self.file_path)
+        from unstructured.partition.text import partition_text
+        return partition_text(filename=txt_file_path, **self.unstructured_kwargs)

requirements.txt

@@ -1,3 +1,6 @@
+pymupdf
+paddlepaddle==2.4.2
+paddleocr
 langchain==0.0.146
 transformers==4.27.1
 unstructured[local-inference]

test_image.py

@@ -0,0 +1,12 @@
+from configs.model_config import *
+import nltk
+
+nltk.data.path = [NLTK_DATA_PATH] + nltk.data.path
+
+filepath = "./img/test.jpg"
+from loader import UnstructuredPaddleImageLoader
+
+loader = UnstructuredPaddleImageLoader(filepath, mode="elements")
+docs = loader.load()
+for doc in docs:
+    print(doc)

test_pdf.py

@@ -0,0 +1,12 @@
+from configs.model_config import *
+import nltk
+
+nltk.data.path = [NLTK_DATA_PATH] + nltk.data.path
+
+filepath = "docs/test.pdf"
+from loader import UnstructuredPaddlePDFLoader
+
+loader = UnstructuredPaddlePDFLoader(filepath, mode="elements")
+docs = loader.load()
+for doc in docs:
+    print(doc)